refactor to follow tool validation

.gitignore

@@ -0,0 +1 @@
+.venv

tool.py

@@ -1,141 +1,10 @@
 import os
 
-from dataclasses import dataclass
-from typing import List, Optional, Tuple
-
-import torch
-from torch.utils.data import DataLoader, Dataset
-
-import base64
-from io import BytesIO
-from PIL import Image
-from pdf2image import convert_from_path
-
-from tqdm import tqdm
-from pqdm.processes import pqdm
-
-from colpali_engine.models import ColQwen2, ColQwen2Processor
-
-from smolagents import Tool, ChatMessage
+from smolagents import Tool
 
 from dotenv import load_dotenv
 load_dotenv()
 
-def encode_image_to_base64(image):
-    """Encodes a PIL image to a base64 string."""
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-DEFAULT_SYSTEM_PROMPT = \
-"""You are a smart assistant designed to answer questions about a PDF document.
-    You are given relevant information in the form of PDF pages preceded by their metadata: document title, page identifier, surrounding context. 
-    Use them to construct a short response to the question, and cite your sources in the following format: (document, page number).
-    If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
-    Give detailed and extensive answers, only containing info in the pages you are given.
-    You can answer using information contained in plots and figures if necessary.
-    Answer in the same language as the query."""
-
-def _build_query(query, pages):
-    messages = []
-    messages.append({"type": "text", "text": "PDF pages:\n"})
-    for page in pages:
-        capt = page.caption
-        if capt is not None:
-            messages.append({
-                    "type": "text",
-                    "text": capt
-                })
-        messages.append({
-                "type": "image_url",
-                "image_url": {
-                "url": f"data:image/jpeg;base64,{encode_image_to_base64(page.image)}"
-                },
-            })
-    messages.append({"type": "text", "text": f"Query:\n{query}"})
-
-    return messages
-
-def query_openai(query, pages, api_key=None, system_prompt=DEFAULT_SYSTEM_PROMPT, model="gpt-4o-mini") -> ChatMessage:
-    """Calls OpenAI's GPT-4o-mini with the query and image data."""
-    if api_key and api_key.startswith("sk"):
-        try:
-            from openai import OpenAI
-    
-            client = OpenAI(api_key=api_key.strip())
-
-            response = client.chat.completions.create(
-                model=model,
-                messages=[
-                    {
-                        "role": "system",
-                        "content": system_prompt
-                    },
-                    {
-                    "role": "user",
-                    "content": _build_query(query, pages)
-                    }
-                ],
-                max_tokens=500,
-            )
-
-            message = ChatMessage.from_dict(
-                response.choices[0].message.model_dump(include={"role", "content", "tool_calls"})
-            )
-            message.raw = response
-
-            return message
-        
-        except Exception as e:
-            return "OpenAI API connection failure. Verify the provided key is correct (sk-***)."
-        
-    return "Enter your OpenAI API key to get a custom response"
-
-CONTEXT_SYSTEM_PROMPT = \
-"""You are a smart assistant designed to extract context of PDF pages.
-    Give concise answers, only containing info in the pages you are given.
-    You can answer using information contained in plots and figures if necessary."""
-
-RAG_SYSTEM_PROMPT = \
-""" You are a smart assistant designed to answer questions about a PDF document.
-
-    You are given relevant information in the form of PDF pages preceded by their metadata: document title, page identifier, surrounding context.
-    Use them to construct a response to the question, and cite your sources.
-    Use the following citation format:
-    "Some information from a first document [1, p.Page Number]. Some information from the same first document but at a different page [1, p.Page Number]. Some more information from another document [2, p.Page Number].
-    ...
-    Sources:
-    [1] Document Title
-    [2] Another Document Title"
-
-    You can answer using information contained in plots and figures if necessary.
-    If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
-    Give detailed answers, only containing info in the pages you are given.
-    Answer in the same language as the query."""
-
-@dataclass
-class Metadata:
-    doc_title: str
-    page_id: int
-    context: Optional[str] = None
-
-    def __str__(self):
-        return f"Document: {self.doc_title}, Page ID: {self.page_id}, Context: {self.context}"
-    
-@dataclass
-class Page:
-    image: Image.Image
-    metadata: Optional[Metadata] = None
-
-    @property
-    def caption(self):
-        if self.metadata is None:
-            return None
-        return f"Document: {self.metadata.doc_title}, Context: {self.metadata.context}"
-    
-    def __hash__(self):
-        return hash(self.image)
-
 class VisualRAGTool(Tool):
     name = "visual_rag"
     description = """Performs a RAG query on your internal PDF documents and returns the generated text response."""
@@ -161,7 +30,13 @@ class VisualRAGTool(Tool):
     model_name: str = "vidore/colqwen2-v1.0"
     api_key: str = os.getenv("OPENAI_KEY")
 
+    def __init__(self, *args, **kwargs):
+        self.is_initialized = False
+
     def _init_models(self, model_name: str) -> None:
+        import torch
+        from colpali_engine.models import ColQwen2, ColQwen2Processor
+
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model = ColQwen2.from_pretrained(
                 model_name,
@@ -170,9 +45,6 @@ class VisualRAGTool(Tool):
                 attn_implementation="flash_attention_2"
             ).eval()
         self.processor = ColQwen2Processor.from_pretrained(model_name)
-
-    def __init__(self, *args, **kwargs):
-        self.is_initialized = False
     
     def setup(self):
         """
@@ -186,8 +58,10 @@ class VisualRAGTool(Tool):
 
         self.is_initialized = True
 
-    def _extract_contexts(self, images, api_key, window=10) -> List[str]:
+    def _extract_contexts(self, images, api_key, window=10) -> list:
         """Extracts context from images."""
+        from utils import query_openai, Page, CONTEXT_SYSTEM_PROMPT
+        from pqdm.processes import pqdm
         try:
             args = [
                 {
@@ -218,8 +92,11 @@ class VisualRAGTool(Tool):
         
         return contexts
 
-    def _preprocess_file(self, file: str, contextualize: bool = True, api_key: str = None, window: int = 10) -> List[Page]:
+    def _preprocess_file(self, file: str, contextualize: bool = True, api_key: str = None, window: int = 10) -> list:
         """Converts a file to images and extracts metadata."""
+        from pdf2image import convert_from_path
+        from utils import Metadata, Page
+
         title = file.split("/")[-1]
         images = convert_from_path(file, thread_count=4)
         if contextualize and api_key:
@@ -230,7 +107,7 @@ class VisualRAGTool(Tool):
 
         return [Page(image=img, metadata=metadata) for img, metadata in zip(images, metadatas)]
 
-    def preprocess(self, files: List[str], contextualize: bool = True, api_key: str = None, window: int = 10) -> List[Page]:
+    def preprocess(self, files: list, contextualize: bool = True, api_key: str = None, window: int = 10) -> list:
         """Preprocesses the files and extracts metadata."""
         pages = [page for file in files for page in self._preprocess_file(file, contextualize=contextualize, api_key=api_key, window=window)]     
 
@@ -238,9 +115,13 @@ class VisualRAGTool(Tool):
 
         return pages
     
-    def compute_embeddings(self, pages: List[Page]) -> List[torch.Tensor]:
+    def compute_embeddings(self, pages) -> list:
         """Embeds the images using the model."""
         """Example script to run inference with ColPali (ColQwen2)"""
+        import torch
+        from torch.utils.data import DataLoader
+        from tqdm import tqdm
+
         # run inference - docs
         dataloader = DataLoader(
             pages,
@@ -259,7 +140,8 @@ class VisualRAGTool(Tool):
 
         return embds
 
-    def index(self, files: List[str], contextualize: bool = True, api_key: str = None, overwrite_db: bool = False) -> int:
+    def index(self, files: list, contextualize: bool = True, api_key: str = None, overwrite_db: bool = False) -> int:
+        """Indexes the uploaded files."""
         if not self.is_initialized:
             self.setup()
         
@@ -285,8 +167,9 @@ class VisualRAGTool(Tool):
         
         return len(embds)
 
-    def retrieve(self, query: str, k: int) -> List[Page]:
+    def retrieve(self, query: str, k: int) -> list:
         """Retrieve the top k documents based on the query."""
+        import torch
         k = min(k, len(self.embds))
 
         qs = []
@@ -307,11 +190,14 @@ class VisualRAGTool(Tool):
 
         return results
         
-    def generate_answer(self, query: str, docs: List[Page], api_key: str = None) -> ChatMessage:
+    def generate_answer(self, query: str, docs: list, api_key: str = None):
+        """Generates an answer based on the query and the retrieved documents."""
+        from utils import query_openai, RAG_SYSTEM_PROMPT
         result = query_openai(query, docs, api_key or self.api_key, system_prompt=RAG_SYSTEM_PROMPT)
         return result
 
-    def search(self, query: str, k: int = 1, api_key: str = None) -> Tuple[list, str]:
+    def search(self, query: str, k: int = 1, api_key: str = None) -> tuple:
+        """Searches for the most relevant pages based on the query."""
         print(f"Searching for query: {query}")
     
         # Retrieve the top k documents
@@ -334,11 +220,9 @@ class VisualRAGTool(Tool):
         #     _ = self.index(files, api_key)
 
         # Retrieve the top k documents and generate response
-        _, rag_answer = self.search(
+        return self.search(
             query=query, 
             files=None, 
             k=k, 
             api_key=api_key
-        )
-
-        return rag_answer
+        )[1]

utils.py

@@ -0,0 +1,124 @@
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import base64
+from io import BytesIO
+from PIL import Image
+
+
+from smolagents import ChatMessage
+
+def encode_image_to_base64(image):
+    """Encodes a PIL image to a base64 string."""
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+DEFAULT_SYSTEM_PROMPT = \
+"""You are a smart assistant designed to answer questions about a PDF document.
+    You are given relevant information in the form of PDF pages preceded by their metadata: document title, page identifier, surrounding context. 
+    Use them to construct a short response to the question, and cite your sources in the following format: (document, page number).
+    If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
+    Give detailed and extensive answers, only containing info in the pages you are given.
+    You can answer using information contained in plots and figures if necessary.
+    Answer in the same language as the query."""
+
+def _build_query(query, pages):
+    messages = []
+    messages.append({"type": "text", "text": "PDF pages:\n"})
+    for page in pages:
+        capt = page.caption
+        if capt is not None:
+            messages.append({
+                    "type": "text",
+                    "text": capt
+                })
+        messages.append({
+                "type": "image_url",
+                "image_url": {
+                "url": f"data:image/jpeg;base64,{encode_image_to_base64(page.image)}"
+                },
+            })
+    messages.append({"type": "text", "text": f"Query:\n{query}"})
+
+    return messages
+
+def query_openai(query, pages, api_key=None, system_prompt=DEFAULT_SYSTEM_PROMPT, model="gpt-4o-mini") -> ChatMessage:
+    """Calls OpenAI's GPT-4o-mini with the query and image data."""
+    if api_key and api_key.startswith("sk"):
+        try:
+            from openai import OpenAI
+    
+            client = OpenAI(api_key=api_key.strip())
+
+            response = client.chat.completions.create(
+                model=model,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": system_prompt
+                    },
+                    {
+                    "role": "user",
+                    "content": _build_query(query, pages)
+                    }
+                ],
+                max_tokens=500,
+            )
+
+            message = ChatMessage.from_dict(
+                response.choices[0].message.model_dump(include={"role", "content", "tool_calls"})
+            )
+            message.raw = response
+
+            return message
+        
+        except Exception as e:
+            return "OpenAI API connection failure. Verify the provided key is correct (sk-***)."
+        
+    return "Enter your OpenAI API key to get a custom response"
+
+CONTEXT_SYSTEM_PROMPT = \
+"""You are a smart assistant designed to extract context of PDF pages.
+    Give concise answers, only containing info in the pages you are given.
+    You can answer using information contained in plots and figures if necessary."""
+
+RAG_SYSTEM_PROMPT = \
+""" You are a smart assistant designed to answer questions about a PDF document.
+
+    You are given relevant information in the form of PDF pages preceded by their metadata: document title, page identifier, surrounding context.
+    Use them to construct a response to the question, and cite your sources.
+    Use the following citation format:
+    "Some information from a first document [1, p.Page Number]. Some information from the same first document but at a different page [1, p.Page Number]. Some more information from another document [2, p.Page Number].
+    ...
+    Sources:
+    [1] Document Title
+    [2] Another Document Title"
+
+    You can answer using information contained in plots and figures if necessary.
+    If it is not possible to answer using the provided pages, do not attempt to provide an answer and simply say the answer is not present within the documents.
+    Give detailed answers, only containing info in the pages you are given.
+    Answer in the same language as the query."""
+
+@dataclass
+class Metadata:
+    doc_title: str
+    page_id: int
+    context: Optional[str] = None
+
+    def __str__(self):
+        return f"Document: {self.doc_title}, Page ID: {self.page_id}, Context: {self.context}"
+    
+@dataclass
+class Page:
+    image: Image.Image
+    metadata: Optional[Metadata] = None
+
+    @property
+    def caption(self):
+        if self.metadata is None:
+            return None
+        return f"Document: {self.metadata.doc_title}, Context: {self.metadata.context}"
+    
+    def __hash__(self):
+        return hash(self.image)