IndustrialKnowledgeAgent: The Smart Industrial Equipment Knowledge Agent

# Core libraries
import os
import json
import functools
import warnings
from typing import List, Dict, Any, Tuple

# LLM and Data Processing
from mistralai 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)

os.environ["MISTRAL_API_KEY"] = "<YOUR MISTRAL API KEY>" # Get your Mistral API key from https://console.mistral.ai/api-keys/

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/

# 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"""

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

# 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

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"],
            },
        },
    }
]

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)))