IndustrialKnowledgeAgent: 스마트 산업 장비 지식 에이전트
IndustrialKnowledgeAgent: 스마트 산업 장비 지식 에이전트 (The Smart Industrial Equipment Knowledge Agent)
산업 환경에서 기술자들은 각종 장비에 대한 포괄적인 정보를 관리하고 검색하는 데 어려움을 겪어요. 이런 정보는 기술 매뉴얼, 정비 기록, 안전 프로토콜, 문제 해결 가이드, 부품 인벤토리에 흩어져 있죠. 이 쿡북에서는 RAG(검색 증강 생성)와 데이터베이스 질의(FunctionCalling)를 통합한 에이전트 워크플로우로 이 문제를 해결하는 방법을 보여줘요.
출처: 문서
본문
문제 정의 (Problem Statement)
산업 현장에서 엔지니어와 기술자들은 다양한 장비에 대한 포괄적인 정보를 관리하고 검색하는 데 종종 어려움을 겪어요. 이 정보는 기술 매뉴얼, 정비 기록, 안전 프로토콜, 문제 해결 가이드, 부품 인벤토리 등에 흩어져 있어요. 데이터가 분산되어 있는 탓에 효과적으로 접근하고 활용하기 어렵고, 이는 비효율성과 잠재적 안전 위험으로 이어져요. 이 문제는 쿼리에 실시간으로 상황에 맞는 응답을 제공하는 지능적이고 적응형인 솔루션을 필요로 해요.
제안된 솔루션 (Proposed Solution)
이런 과제를 해결하기 위해, RAG(Retrieval-Augmented Generation) 시스템과 **데이터베이스 질의 시스템(FunctionCalling)**을 통합한 에이전트 워크플로우를 제안해요. 이 솔루션은 LLM(구조화된 출력 메커니즘 포함), 임베딩 모델, 구조화된 데이터 검색을 활용해서 상황에 맞고 정확한 정보를 제공해요. 워크플로우는 각각 특정 역할을 가진 여러 에이전트가 오케스트레이션해요.
- RAGAgent: LLM과 임베딩 모델을 사용해서 기술 문서에서 상황에 맞는 관련 정보를 검색하고 생성해요.
- DatabaseQueryAgent: 정비 기록, 기술 사양, 부품 인벤토리, 규정 준수 기록이 담긴 데이터베이스에서 정확하고 구조화된 데이터를 검색해요.
- WorkflowOrchestrator: RAGSearchAgent와 DatabaseAgent 간의 상호작용을 조율해서 매끄럽고 효율적인 쿼리 해결을 보장해요.
데이터셋 세부사항 (Dataset Details)
PDF 문서 (PDF Documents)
PDF 문서에는 다양한 산업 장비에 대한 상세 정보가 담겨 있으며, 다음과 같이 분류돼요.
- 기술 매뉴얼 (Technical Manuals): 작동 및 유지보수 가이드
- 정비 가이드 (Maintenance Guides): 일상 및 예방 정비 작업
- 문제 해결 가이드 (Troubleshooting Guides): 일반적인 문제에 대한 해결책
- 안전 프로토콜 (Safety Protocols): 안전 절차 및 지침
데이터베이스 (Databases)
데이터베이스에는 PDF 문서를 보완하는 구조화된 정보가 들어 있어요.
- 규정 준수 데이터베이스 (
compliance_db): 안전 인증 및 규정 준수 상태 - 정비 데이터베이스 (
maintenance_db): 정비 활동 기록 - 기술 사양 데이터베이스 (
technical_specifications_db): 상세한 기술 사양 - 부품 인벤토리·호환성 데이터베이스 (
parts_inventory_compatibility_db): 부품, 호환성, 재고 상태 정보
이 데이터셋들을 통합함으로써, 제안된 에이전트 워크플로우는 산업 장비 정보를 관리하고 검색하는 포괄적이고 효율적인 시스템을 제공해서, 엔지니어와 기술자가 가장 관련성 높고 최신의 정보에 접근할 수 있도록 돕는 것을 목표로 해요.
참고(NOTE): 이 데모에서 사용된 모든 데이터는 합성으로 생성되었어요.
기술 아키텍처 (Technical Architecture)
설치 (Installation)
IndusAgent 시스템에 필요한 Python 패키지를 설치해요.
Python
!pip install mistralai==1.5.1 # Mistral AI client
!pip install qdrant-client==1.13.2 # Vector database client
!pip install gdown==5.2.0 # Google Drive download
임포트 (Imports)
LLM 연산, 데이터 처리, 벡터 데이터베이스 관리, 유틸리티 함수에 필요한 라이브러리를 가져와요.
# Core libraries
import os
import json
import functools
import warnings
from typing import List, Dict, Any, Tuple
# LLM and Data Processing
from mistralai.client import Mistral
from pydantic import BaseModel
import pandas as pd
import sqlite3
from tqdm import tqdm
# Vector Database
from qdrant_client import QdrantClient
from qdrant_client.models import (
PointStruct, VectorParams, Distance,
Filter, FieldCondition, MatchValue
)
# Data Download
import gdown
import zipfile
# Suppress warnings
warnings.filterwarnings('ignore', category=DeprecationWarning)
데이터 다운로드 (Download Data)
Google Drive에서 데이터셋을 다운로드해서 데이터 디렉터리에 풀고 작업 환경을 설정해요. 데이터셋에는 데이터베이스 연산용 CSV 파일과 문서 처리를 위한 PDF가 포함돼 있어요.
Google Drive에서 데이터 다운로드 (Download data from Google Drive)
Python Output
file_id = "1lwYSN6ry3JOA7pw3WAx72a_IXGqqmR8y"
output_file = "data.zip" # Change this if your file is not a ZIP file
# Google Drive direct download URL
gdrive_url = f"https://drive.google.com/uc?id={file_id}"
# Download the file
gdown.download(gdrive_url, output_file, quiet=False)
print(f"✅ File downloaded: {output_file}")
데이터 디렉터리 추출 및 설정 (Extract and setup data directory)
Python Output
# Unzip the file into the current directory
with zipfile.ZipFile(output_file, 'r') as zip_ref:
zip_ref.extractall(".") # Extracts directly to the current directory
print(f"✅ Files extracted to: {os.getcwd()}") # Confirm extraction path
output_dir = "data"
# Change working directory to the extracted folder
os.chdir(output_dir)
# Verify the new working directory
print(f"📂 Current directory: {os.getcwd()}")
Python
# List files in the extracted folder
print("📜 Extracted files:", os.listdir())
환경 변수 설정 (Set up environment variables)
인증을 위해 Mistral API 키를 환경 변수로 설정해요.
os.environ["MISTRAL_API_KEY"] = "<YOUR MISTRAL API KEY>" # Get your Mistral API key from https://console.mistral.ai/api-keys/
Mistral LLM 및 Qdrant 벡터 데이터베이스 초기화 (Initialize Mistral LLM and Qdrant Vector Database)
텍스트 생성을 위한 Mistral LLM 클라이언트와 유사도 검색 연산을 위한 Qdrant 벡터 데이터베이스 클라이언트를 초기화해요.
참고(Note):
데모에는 최신 모델인 Mistral Small 3을 사용할 거예요.
진행 전에 Qdrant Cloud나 Docker를 설정해야 해요. 설정 지침은 문서에서 확인할 수 있어요.
model = "mistral-small-latest"
mistral_client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
qdrant_client = QdrantClient(
url= "<URL>",
api_key= "<API KEY>",
) # Replace with your Qdrant API key and URL if you are using Qdrant Cloud - https://cloud.qdrant.io/
시스템 프롬프트 (System Prompts)
시스템은 응답 생성을 위해 LLM을 안내하는 세 가지 유형의 프롬프트를 사용해요.
- PDF 요약 프롬프트 (
summarization_prompt): PDF 문서의 간결한 요약을 만드는 데 사용해요. - 응답 생성 프롬프트 (
response_generation_prompt): 검색된 컨텍스트를 바탕으로 응답을 생성하는 데 사용해요. - 최종 응답 통합 프롬프트 (
final_response_generation_prompt): PDF와 여러 데이터베이스 등 여러 소스의 응답을 요약하는 데 사용해요.
# Define the prompt for generating a response
response_generation_prompt = '''Based on the following context answer the query:\n\n Context: {context}\n\n Query: {query}'''
# Prompt for summarizing the PDF text
summarization_prompt = '''Your task is to summarize the following text focusing on the core essence of the text in maximum of 2-3 sentences.'''
# Prompt for final response summarization
final_response_summarization_prompt = """You are an expert technical assistant. Your task is to create a comprehensive,
coherent response by combining information from multiple sources: database records and documentation.
Consider the following guidelines:
1. Integrate information from both sources seamlessly
2. Resolve any conflicts between sources, if they exist
3. Present information in a logical, step-by-step manner when applicable
4. Include specific technical details, measurements, and procedures when available
5. Prioritize safety-related information when present
6. Add relevant maintenance intervals or schedules if mentioned
7. Reference specific part numbers or specifications when provided
The user's query is: {query}
Based on the following responses from different sources, create a unified, clear answer:
{responses}
Remember to:
- Focus on accuracy and completeness
- Maintain technical precision
- Use clear, professional language
- Address all aspects of the query
- Highlight any important warnings or precautions"""
DataProcessor
DataProcessor 클래스는 시스템의 모든 데이터 처리 연산을 담당하는 포괄적인 컴포넌트예요. 비정형(PDF) 데이터와 정형(CSV) 데이터를 모두 관리하고, 임베딩 생성과 저장도 처리해요.
- Mistral OCR을 사용한 PDF 문서 처리와 텍스트 추출
- CSV를 데이터베이스로 변환
- 임베딩 생성 및 벡터 저장
- 문서와 데이터의 배치 처리
주요 컴포넌트 (Main Components)
1. 문서 처리 (Document Processing)
get_categorized_filepaths: 디렉터리 구조를 탐색해서 카테고리별 PDF 파일 경로를 가져와요.parse_pdf: Mistral OCR을 사용해 PDF 파일의 모든 페이지에서 텍스트를 추출해요.process_single_pdf: 개별 PDF를 전체 파이프라인으로 처리해요.process_documents: 여러 문서를 순차적으로 처리해요.
2. 요약과 임베딩 (Summarization and Embeddings)
summarize: Mistral 모델을 사용해 텍스트의 간결한 요약을 생성해요.get_text_embedding: Mistral의 임베딩 모델로 텍스트 임베딩을 생성해요.qdrant_insert_embeddings: 메타데이터와 함께 임베딩을 Qdrant 벡터 데이터베이스에 저장해요.process_and_store_embeddings: 임베딩의 배치 처리를 담당해요.
3. 데이터베이스 연산 (Database Operations)
insert_csv_to_table: 단일 CSV 파일을 지정된 데이터베이스 테이블에 로드해요.insert_data_database: 여러 CSV 파일을 각각의 테이블에 삽입해요.
class DataProcessor:
"""
Handles all data processing operations including:
- PDF parsing and text extraction
- CSV to database ingestion
- Embedding generation and storage
- Batch processing of documents and data
"""
def __init__(self, mistral_client: Mistral, qdrant_client: QdrantClient):
self.mistral_client = mistral_client
self.qdrant_client = qdrant_client
def get_categorized_filepaths(self, root_dir: str) -> List[Dict[str, str]]:
"""
Walk through the directory structure and get file paths with their categories.
"""
categorized_files = []
for category in os.listdir(root_dir):
category_path = os.path.join(root_dir, category)
if not os.path.isdir(category_path):
continue
for root, _, files in os.walk(category_path):
for file in files:
if file.lower().endswith('.pdf'):
filepath = os.path.join(root, file)
categorized_files.append({
'filepath': filepath,
'category': category
})
return categorized_files
def parse_pdf(self, file_path: str) -> str:
"""Parse a PDF file and extract text from all pages using Mistral OCR."""
# Upload a file
uploaded_pdf = self.mistral_client.files.upload(
file={
"file_name": file_path,
"content": open(file_path, "rb"),
},
purpose="ocr"
)
# Get a signed URL for the uploaded file
signed_url = self.mistral_client.files.get_signed_url(file_id=uploaded_pdf.id)
# Get OCR results
ocr_response = self.mistral_client.ocr.process(
model="mistral-ocr-latest",
document={
"type": "document_url",
"document_url": signed_url.url,
}
)
# Extract text from the OCR response
text = "\n".join([x.markdown for x in (ocr_response.pages)])
return text
def summarize(self, text: str, summarization_prompt: str = summarization_prompt) -> str:
"""Summarize the given text using the Mistral model."""
chat_response = self.mistral_client.chat.complete(
model=model,
messages=[
{
"role": "system",
"content": summarization_prompt
},
{
"role": "user",
"content": text
},
],
temperature=0
)
return chat_response.choices[0].message.content
def get_text_embedding(self, inputs: List[str]) -> List[float]:
"""Get the text embedding for the given inputs."""
embeddings_batch_response = self.mistral_client.embeddings.create(
model="mistral-embed",
inputs=inputs
)
return embeddings_batch_response.data[0].embedding
def qdrant_insert_embeddings(self, summaries: List[str], texts: List[str],
filepaths: List[str], categories: List[str]):
"""Insert embeddings into Qdrant with metadata."""
embeddings = [self.get_text_embedding([t]) for t in summaries]
if not self.qdrant_client.collection_exists("embeddings"):
self.qdrant_client.create_collection(
collection_name="embeddings",
vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
)
self.qdrant_client.upsert(
collection_name="embeddings",
points=[
PointStruct(
id=idx,
vector=embedding,
payload={
"filepath": filepaths[idx],
"category": categories[idx],
"text": texts[idx]
}
) for idx, embedding in enumerate(embeddings)
]
)
def process_single_pdf(self, file_info: Dict[str, str]) -> Dict[str, any]:
"""Process a single PDF file through the pipeline."""
filepath = file_info['filepath']
category = file_info['category']
pdf_text = self.parse_pdf(filepath)
summary = self.summarize(pdf_text)
return {
'filepath': filepath,
'category': category,
'full_text': pdf_text,
'summary': summary
}
def process_documents(self, file_list: List[Dict[str, str]]) -> List[Dict[str, any]]:
"""Process documents sequentially."""
processed_docs = []
for file_info in tqdm(file_list, desc="Processing PDFs"):
try:
processed_doc = self.process_single_pdf(file_info)
processed_docs.append(processed_doc)
except Exception as e:
print(f"Error processing {file_info['filepath']}: {str(e)}")
continue
return processed_docs
def insert_csv_to_table(self, file_path: str, db_path: str, table_name: str):
"""
Insert CSV data into a table of SQLite database.
Args:
file_path (str): Path to the CSV file
db_path (str): Path to the SQLite database
table_name (str): Name of the table to create/update
"""
df = pd.read_csv(file_path)
conn = sqlite3.connect(db_path)
df.to_sql(table_name, conn, if_exists='replace', index=False)
conn.close()
def insert_data_database(self, db_path: str, file_mappings: Dict[str, str]):
"""
Bulk insert multiple CSV files into their respective database tables.
Args:
db_path (str): Path to the SQLite database
file_mappings (Dict[str, str]): Dictionary mapping table names to CSV file paths
"""
for table_name, file_path in file_mappings.items():
try:
self.insert_csv_to_table(file_path, db_path, table_name)
print(f"Successfully inserted data into {table_name}")
except Exception as e:
print(f"Error inserting data into {table_name}: {str(e)}")
def process_and_store_embeddings(self, docs: List[Dict[str, any]], batch_size: int = 10):
"""Generate embeddings and store them in Qdrant in batches."""
for i in range(0, len(docs), batch_size):
batch = docs[i:i + batch_size]
texts = [doc['full_text'] for doc in batch]
summaries = [doc['summary'] for doc in batch]
filepaths = [doc['filepath'] for doc in batch]
categories = [doc['category'] for doc in batch]
try:
self.qdrant_insert_embeddings(summaries, texts, filepaths, categories)
print(f"Processed batch {i//batch_size + 1}/{(len(docs) + batch_size - 1)//batch_size}")
except Exception as e:
print(f"Error processing batch starting at index {i}: {str(e)}")
continue
RAGAgent
RAGAgent 클래스는 지능적인 검색과 응답 생성을 위한 RAG(검색 증강 생성)를 구현해요. 벡터 검색 기능과 LLM을 결합해서 상황에 맞는 답변을 제공해요.
- 쿼리 분류 및 카테고리화
- Qdrant에서의 벡터 유사도 검색
- 컨텍스트 인지 응답 생성
- 문서 인용 처리
주요 컴포넌트 (Main Components)
1. 쿼리 처리 (Query Processing)
query_categorization: 쿼리를 미리 정의된 카테고리(기술 매뉴얼, 안전 프로토콜 등)로 분류해요.query: 쿼리부터 최종 응답까지 전체 RAG 파이프라인을 오케스트레이션해요.
2. 검색 및 검색결과 (Search and Retrieval)
qdrant_search: 쿼리 임베딩으로 의미론적 검색을 수행하고, 문서 카테고리로 결과를 필터링해서 상위 k개의 관련 문서를 반환해요.
3. 응답 생성 (Response Generation)
generate_response: 검색된 컨텍스트로 자연어 응답을 만들고, 전문화된 프롬프트로 LLM을 사용하며, 소스 문서에 대한 인용을 제공해요.
쿼리 카테고리 모델 (Query Category Model)
RAGAgent가 쿼리를 관련 카테고리(technical_manual, safety_protocol 등)로 분류하는 데 사용하는 구조를 정의하는 Pydantic 모델이에요.
# Define category model for query classification
class Category(BaseModel):
category: str
class RAGAgent:
"""
Agent responsible for Retrieval-Augmented Generation (RAG) operations.
"""
def __init__(self, mistral_client: Mistral, qdrant_client: QdrantClient):
self.mistral_client = mistral_client
self.qdrant_client = qdrant_client
def generate_response(self, context: str, query: str) -> str:
"""Generate a response based on the given context and query."""
chat_response = self.mistral_client.chat.complete(
model=model,
messages=[
{
"role": "user",
"content": response_generation_prompt.format(context=context, query=query)
},
]
)
return chat_response.choices[0].message.content
def query_categorization(self, query: str) -> str:
"""Categorize the query into predefined categories."""
chat_response = self.mistral_client.chat.parse(
model=model,
messages=[
{
"role": "system",
"content": "Classify the query into one or more categories of the following list: ['technical_manual', 'safety_protocol', 'maintenance_guide', 'troubleshooting_guide']"
},
{
"role": "user",
"content": query,
},
],
response_format=Category,
max_tokens=256,
temperature=0
)
return json.loads(chat_response.choices[0].message.content)
def qdrant_search(self, query: str, category: str = None, top_k: int = 5) -> List[Dict[str, Any]]:
"""Search for similar texts in Qdrant based on the query and category."""
query_vector = DataProcessor(self.mistral_client, self.qdrant_client).get_text_embedding([query])
retrieval_results = self.qdrant_client.search(
collection_name="embeddings",
query_vector=query_vector,
query_filter=Filter(
must=[
FieldCondition(
key='category',
match=MatchValue(value=category)
)
]
),
limit=top_k
)
return retrieval_results
def query(self, query_text: str, top_k: int = 3) -> Tuple[str, str]:
"""Process a natural language query using RAG."""
category = self.query_categorization(query_text)["category"]
results = self.qdrant_search(query_text, category, top_k=top_k)
file_paths = [result.payload["filepath"] for result in results]
retrieved_text = "\n".join([result.payload["text"] for result in results])
citations = ",".join([result.payload["filepath"] for result in results])
return self.generate_response(retrieved_text, query_text), citations
데이터베이스 쿼리 도구 (Database Query Tools)
데이터베이스 쿼리 함수 도구를 정의해요. 이 도구들은 DatabaseQueryAgent의 함수 호출 인터페이스를 정의해서, 서로 다른 데이터베이스 테이블을 구조적으로 쿼리할 수 있게 해줘요.
tools = [
{
"type": "function",
"function": {
"name": "query_compliance",
"description": '''Query compliance records with filters. \n\n A sample example of columns and corresponding values from db are:\n\n EquipmentID,EquipmentName,Manufacturer,Model,ComplianceType,Certification,IssueDate,ExpiryDate,ComplianceStatus,ResponsiblePerson
1,CNC Machine,ABC Corp,Model X,Safety,ISO 9001,2020-01-15,2025-01-15,Active,John Doe''',
"parameters": {
"type": "object",
"properties": {
"filters": {
"type": "object",
"description": '''Dictionary of column names and values to filter by.''',
"additionalProperties": {
"type": "string"
}
}
},
"required": ["filters"],
},
},
},
{
"type": "function",
"function": {
"name": "query_maintenance",
"description": '''Query maintenance records with filters. \n\n A sample example of columns and corresponding values from db are:\n\n EquipmentID,EquipmentName,Manufacturer,Model,InstallationDate,LastMaintenanceDate,NextMaintenanceDate,MaintenanceType,MaintenanceDetails,MaintenanceStatus,ResponsibleTechnician
1,CNC Machine,ABC Corp,Model X,2020-01-15,2023-09-01,2023-12-01,Preventive,Oil change,Completed,John Doe''',
"parameters": {
"type": "object",
"properties": {
"filters": {
"type": "object",
"description": "Dictionary of column names and values to filter by",
"additionalProperties": {
"type": "string"
}
}
},
"required": ["filters"],
},
},
},
{
"type": "function",
"function": {
"name": "query_technical_specs",
"description": '''Query technical specifications with filters.\n\n A sample example of columns and corresponding values from db are:\n\n EquipmentID,EquipmentName,Manufacturer,Model,SpecificationType,SpecificationDetail,Unit,Value,DateMeasured,MeasuredBy
1,CNC Machine,ABC Corp,Model X,Power,Motor Power,kW,15,2023-01-15,John Doe''',
"parameters": {
"type": "object",
"properties": {
"filters": {
"type": "object",
"description": "Dictionary of column names and values to filter by",
"additionalProperties": {
"type": "string"
}
}
},
"required": ["filters"],
},
},
},
{
"type": "function",
"function": {
"name": "query_parts_inventory_compatibility",
"description": '''Query parts, inventory and compatibility with filters.\n\n A sample example of columns and corresponding values from db are:\n\n PartID,PartName,EquipmentID,EquipmentName,Manufacturer,Model,PartType,Quantity,Compatibility,Supplier,LastOrderDate,NextOrderDate,PartStatus
1,Oil Filter,1,CNC Machine,ABC Corp,Model X,Filter,50,Compatible,Supplier A,2023-01-15,2023-12-01,In Stock''',
"parameters": {
"type": "object",
"properties": {
"filters": {
"type": "object",
"description": "Dictionary of column names and values to filter by",
"additionalProperties": {
"type": "string"
}
}
},
"required": ["filters"],
},
},
}
]
DatabaseQueryAgent
DatabaseQueryAgent 클래스는 SQLite 데이터베이스와의 상호작용을 관리해요. 정비 기록, 기술 사양, 부품 인벤토리, 규정 준수 기록이 담긴 여러 데이터베이스에 대한 함수 호출을 통해 구조화된 데이터 쿼리를 처리해요. 서로 다른 데이터베이스 테이블에 대한 전문화된 쿼리 기능을 제공해요.
- 자연어 쿼리 처리
- 구조화된 데이터베이스 쿼리
- 쿼리 실행을 위한 함수 호출
- JSON 응답 포맷팅
주요 컴포넌트 (Main Components)
1. 테이블별 쿼리 (Table-Specific Queries)
query_compliance: 필터링된 규정 준수 레코드를 검색해요.query_maintenance: 정비 관련 정보에 접근해요.query_technical_specs: 기술 사양을 가져와요.query_parts_inventory_compatibility: 부품 및 호환성 데이터를 검색해요.
2. 쿼리 처리 (Query Processing)
query: 함수 호출을 사용해 자연어 쿼리를 처리하고, 적절한 데이터베이스 연산을 위한 도구 호출을 다루며, 데이터베이스 도구 인용을 추적하고, 쿼리 결과로 응답을 포맷해요.
class DatabaseQueryAgent:
"""
Agent responsible for interacting with the SQLite database.
"""
def __init__(self, db_path: str, mistral_client: Mistral):
self.db_path = db_path
self.tools = tools
self.names_to_functions = {
'query_compliance': functools.partial(self.query_compliance),
'query_maintenance': functools.partial(self.query_maintenance),
'query_technical_specs': functools.partial(self.query_technical_specs),
'query_parts_inventory_compatibility': functools.partial(self.query_parts_inventory_compatibility)
}
self.mistral_client = mistral_client
def query_compliance(self, filters: Dict[str, str]) -> str:
"""
Query compliance table with filters.
Args:
filters (Dict[str, str]): Dictionary of column names and values to filter by.
Returns:
str: The query result in JSON format.
"""
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
where_conditions = []
params = []
for column, value in filters.items():
where_conditions.append(f"{column} = ?")
params.append(value)
where_clause = " AND ".join(where_conditions)
query = f"SELECT * FROM compliance WHERE {where_clause}"
cursor.execute(query, params)
columns = [description[0] for description in cursor.description]
result = cursor.fetchone()
if result:
record = dict(zip(columns, result))
return json.dumps({'result': record})
return json.dumps({'error': 'No matching records found'})
except Exception as e:
return json.dumps({'error': str(e)})
finally:
conn.close()
def query_maintenance(self, filters: Dict[str, str]) -> str:
"""
Query maintenance table with filters.
Args:
filters (Dict[str, str]): Dictionary of column names and values to filter by.
Returns:
str: The query result in JSON format.
"""
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
where_conditions = []
params = []
for column, value in filters.items():
where_conditions.append(f"{column} = ?")
params.append(value)
where_clause = " AND ".join(where_conditions)
query = f"SELECT * FROM maintenance WHERE {where_clause}"
cursor.execute(query, params)
columns = [description[0] for description in cursor.description]
result = cursor.fetchone()
if result:
record = dict(zip(columns, result))
return json.dumps({'result': record})
return json.dumps({'error': 'No matching records found'})
except Exception as e:
return json.dumps({'error': str(e)})
finally:
conn.close()
def query_technical_specs(self, filters: Dict[str, str]) -> str:
"""
Query technical specifications table with filters.
Args:
filters (Dict[str, str]): Dictionary of column names and values to filter by.
Returns:
str: The query result in JSON format.
"""
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
where_conditions = []
params = []
for column, value in filters.items():
where_conditions.append(f"{column} = ?")
params.append(value)
where_clause = " AND ".join(where_conditions)
query = f"SELECT * FROM technical_specifications WHERE {where_clause}"
cursor.execute(query, params)
columns = [description[0] for description in cursor.description]
result = cursor.fetchone()
if result:
record = dict(zip(columns, result))
return json.dumps({'result': record})
return json.dumps({'error': 'No matching records found'})
except Exception as e:
return json.dumps({'error': str(e)})
finally:
conn.close()
def query_parts_inventory_compatibility(self, filters: Dict[str, str]) -> str:
"""
Query parts inventory and compatibility table with filters.
Args:
filters (Dict[str, str]): Dictionary of column names and values to filter by.
Returns:
str: The query result in JSON format.
"""
try:
conn = sqlite3.connect(self.db_path)
cursor = conn.cursor()
where_conditions = []
params = []
for column, value in filters.items():
where_conditions.append(f"{column} = ?")
params.append(value)
where_clause = " AND ".join(where_conditions)
query = f"SELECT * FROM parts_inventory_compatibility WHERE {where_clause}"
cursor.execute(query, params)
columns = [description[0] for description in cursor.description]
result = cursor.fetchone()
if result:
record = dict(zip(columns, result))
return json.dumps({'result': record})
return json.dumps({'error': 'No matching records found'})
except Exception as e:
return json.dumps({'error': str(e)})
finally:
conn.close()
def query(self, query_text: str) -> str:
"""
Process a natural language query using the database tools.
Args:
query_text (str): Natural language query
Returns:
str: Response from the database query
"""
messages = [{"role": "user", "content": query_text}]
# Get initial response with potential tool calls
response = self.mistral_client.chat.complete(
model=model,
messages=messages,
tools=self.tools,
tool_choice="any",
)
messages.append(response.choices[0].message)
citations = set()
# Handle any tool calls
if hasattr(response.choices[0].message, 'tool_calls') and response.choices[0].message.tool_calls:
for tool_call in response.choices[0].message.tool_calls:
function_name = tool_call.function.name
print(f"Tool call: {function_name}")
citations.add(function_name)
function_params = json.loads(tool_call.function.arguments)
print(f"Tool call parameters: {function_params}")
function_result = self.names_to_functions[function_name](**function_params)
messages.append({
"role": "tool",
"name": function_name,
"content": function_result,
"tool_call_id": tool_call.id
})
# Get final response
final_response = self.mistral_client.chat.complete(
model="mistral-small-latest",
messages=messages
)
return final_response.choices[0].message.content, ",".join(list((citations)))
WorkflowOrchestrator
WorkflowOrchestrator 클래스는 RAGAgent와 DatabaseQueryAgent 간의 상호작용을 조율해서, 구조화된 데이터 소스와 비구조화된 데이터 소스의 정보를 결합한 포괄적인 응답을 제공해요.
- 워크플로우 오케스트레이션 및 조정
- 응답 결합과 통합
- 최종 응답 요약
- 소스 인용 관리
주요 컴포넌트 (Main Components)
1. 워크플로우 실행 (Workflow Execution)
workflow: 전체 쿼리 처리 파이프라인을 관리하고, 두 에이전트의 응답을 조정하며, 최종 통합 응답을 생성하고, 인용을 통한 추적성을 유지해요.
2. 응답 요약 (Response Summarization)
combine_and_summarize_responses: 두 에이전트의 응답을 병합하고 요약하며, 결합된 응답에 구조화된 포맷을 적용하고, 일관된 출력을 위해 요약 프롬프트를 사용해요.
class WorkflowOrchestrator:
"""
WorkflowOrchestrator is responsible for orchestrating the workflow between RAGSearchAgent and DatabaseQueryAgent.
Handles query processing, response combination, and final summarization.
"""
def __init__(self,
rag_agent: RAGAgent,
db_query_agent: DatabaseQueryAgent,
client: Mistral):
"""
Initialize WorkflowOrchestrator with necessary components.
Args:
rag_agent: RAGSearchAgent for document retrieval and generation
db_query_agent: DatabaseQueryAgent for structured data queries
client: Mistral client for text generation
"""
self.rag_agent = rag_agent
self.db_query_agent = db_query_agent
self.client = client
def combine_and_summarize_responses(self,
responses: Dict[str, str],
query: str,
summarization_prompt: str = summarization_prompt) -> str:
"""
Combine and summarize multiple responses into a coherent final response.
Args:
responses: Dictionary of response types and their content
query: Original user query
summarization_prompt: Template for summarization
Returns:
str: Summarized and combined response
"""
# Format responses into a structured text
combined_text = "\n\n".join([
f"{source}: {content}"
for source, content in responses.items()
])
# Generate summarized response
chat_response = self.client.chat.complete(
model=model,
messages=[
{
"role": "system",
"content": summarization_prompt
},
{
"role": "user",
"content": f"Query: {query}\n\nResponses:\n{combined_text}"
},
],
temperature=0
)
return chat_response.choices[0].message.content
def workflow(self, query: str) -> str:
"""
Execute the workflow for processing a query.
Args:
query: User query
Returns:
str: Final response with citations
"""
# Get responses from both agents
db_response, tools_citations = self.db_query_agent.query(query)
rag_response, rag_citations = self.rag_agent.query(query)
# Combine responses into a dictionary
responses = {
"Database Response": db_response,
"RAG Response": rag_response
}
# Generate final summarized response
final_response = self.combine_and_summarize_responses(
responses=responses,
query=query,
summarization_prompt=final_response_summarization_prompt
)
# Add citations
citations = (
f"\n\nSources:\n"
f"- Database Tools: {((tools_citations))}\n"
f"- PDF Sources: {rag_citations}"
)
return final_response + citations
초기화 및 문서 처리 (Initialize and Process Documents)
DataProcessor를 초기화하고 파일 수집부터 임베딩 저장까지 전체 파이프라인으로 PDF 문서를 처리해요.
Python
# Initialize the processor
doc_processor = DataProcessor(mistral_client, qdrant_client)
# Process documents
file_list = doc_processor.get_categorized_filepaths(root_dir='./pdf_data')
processed_docs = doc_processor.process_documents(file_list)
doc_processor.process_and_store_embeddings(processed_docs)
데이터베이스 테이블에 데이터 삽입 (Insert Data into Database tables)
여러 CSV 파일을 SQLite의 해당 테이블에 로드해요.
Python
# Insert data into tables
db_path = "./database.db"
file_mappings = {
"compliance": "./csv_data/compliance_db.csv",
"maintenance": "./csv_data/maintenance_db.csv",
"technical_specifications": "./csv_data/technical_specifications_db.csv",
"parts_inventory_compatibility": "./csv_data/parts_inventory_compatibility_db.csv"
}
doc_processor.insert_data_database(db_path, file_mappings)
에이전트 초기화 (Initialise the Agents)
세 가지 핵심 에이전트를 초기화해요.
- 문서 검색과 응답을 위한 RAGAgent
- 구조화된 데이터 쿼리를 위한 DatabaseQueryAgent
- 응답을 오케스트레이션하는 WorkflowAgent
rag_agent = RAGAgent(mistral_client, qdrant_client)
db_query_agent = DatabaseQueryAgent(db_path, mistral_client)
workflow_orchestrator = WorkflowOrchestrator(rag_agent, db_query_agent, mistral_client)
예제 쿼리 (Example Queries)
Python Output
query = "What are the troubleshooting steps for inaccurate machining in CNC Machine (Model X) and when was its last maintenance performed?"
print(f"Query: {query}")
print("----------------------")
answer = workflow_orchestrator.workflow(query)
print("------------Answer----------")
print(answer)
Python
query = "What are the safety protocols for the Cooling System (Model Y), and when is its next scheduled maintenance?"
print(f"Query: {query}")
print("----------------------")
answer = workflow_orchestrator.workflow(query)
print("------------Answer----------")
print(answer)