Fix utils bug

tool.py

@@ -29,6 +29,26 @@ class VisualRAGTool(Tool):
 
     model_name: str = "vidore/colqwen2-v1.0"
     api_key: str = os.getenv("OPENAI_KEY")
+        
+    class Page:
+        from typing import Optional, Dict, Any
+        from PIL import Image
+    
+        image: Image.Image
+        metadata: Optional[Dict[str, Any]] = None
+
+        def __init__(self, image, metadata=None):
+            self.image = image
+            self.metadata = metadata
+
+        @property
+        def caption(self):
+            if self.metadata is None:
+                return None
+            return f"Document: {self.metadata.get('doc_title')}, Context: {self.metadata.get('context')}"
+        
+        def __hash__(self):
+            return hash(self.image)
 
     def __init__(self, *args, **kwargs):
         self.is_initialized = False
@@ -58,20 +78,103 @@ class VisualRAGTool(Tool):
 
         self.is_initialized = True
 
+    def _encode_image_to_base64(self, image):
+        """Encodes a PIL image to a base64 string."""
+        from io import BytesIO
+        import base64
+
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
+    
+    def _build_query(self, query: str, pages: list) -> list:
+        """Builds the query for OpenAI based on the pages and the query."""
+        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,{self._encode_image_to_base64(page.image)}"
+                    },
+                })
+        messages.append({"type": "text", "text": f"Query:\n{query}"})
+
+        return messages
+
+    def query_openai(self, query, pages, api_key=None, system_prompt=None, model="gpt-4o-mini"):
+        """Calls OpenAI's GPT-4o-mini with the query and image data."""
+        from smolagents import ChatMessage
+
+        system_prompt = system_prompt or \
+        """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."""
+        
+        api_key = api_key or self.api_key
+
+        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": self._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"
+
     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
+        from pqdm.threads import pqdm
+
+        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."""
+
         try:
             args = [
                 {
                     'query': "Give the general context about these pages. Give the context in the same language as the documents.",
-                    'pages': [Page(image=im) for im in images[max(i-window+1, 0):i+1]],
+                    'pages': [self.Page(image=im) for im in images[max(i-window+1, 0):i+1]],
                     'api_key': api_key,
                     'system_prompt': CONTEXT_SYSTEM_PROMPT,
                 } for i in range(0, len(images), window)
             ]
-            window_contexts = pqdm(args, query_openai, n_jobs=8, argument_type='kwargs')
+            window_contexts = pqdm(args, self.query_openai, n_jobs=8, argument_type='kwargs')
 
             # code sequentially ftm with tqdm
             # query = "Give the general context about these pages. Give the context in the same language as the documents."
@@ -95,7 +198,6 @@ class VisualRAGTool(Tool):
     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)
@@ -103,15 +205,15 @@ class VisualRAGTool(Tool):
             contexts = self._extract_contexts(images, api_key, window=window)
         else:
             contexts = [None for _ in range(len(images))]
-        metadatas = [Metadata(doc_title=title, page_id=i, context=contexts[i]) for i in range(len(images))]
+        metadatas = [{'doc_title': title, 'page_id': i, 'context': contexts[i]} for i in range(len(images))]
 
-        return [Page(image=img, metadata=metadata) for img, metadata in zip(images, metadatas)]
+        return [self.Page(image=img, metadata=metadata) for img, metadata in zip(images, metadatas)]
 
     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)]     
 
-        print(f"Example metadata:\n{pages[0].metadata.context}")
+        print(f"Example metadata:\n{pages[0].metadata.get('context')}")
 
         return pages
     
@@ -183,7 +285,7 @@ class VisualRAGTool(Tool):
         top_k_idx = scores.topk(k).indices.tolist()
 
         print("Top Scores:")
-        [print(f'Page {self.pages[idx].metadata.page_id}: {scores[idx]}') for idx in top_k_idx]
+        [print(f'Page {self.pages[idx].metadata.get('page_id')}: {scores[idx]}') for idx in top_k_idx]
 
         # Get the top k results
         results = [self.pages[idx] for idx in top_k_idx]
@@ -192,8 +294,25 @@ class VisualRAGTool(Tool):
         
     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)
+
+        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."""
+        
+        result = self.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:

utils.py

@@ -1,124 +0,0 @@
-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)