feat: retrieval (1st version)

.gitignore

@@ -11,9 +11,13 @@ my-traffic-analysis-441217-32bda1474a0f.json
 # Python
 *__pycache__/
 
-# Model
+# Project
 llm/
 attachments/
 logs/
 1st_run/
-metadata.jsonl
+metadata.jsonl
+tests.py
+
+chroma_db/
+*.bin

prompt.md

@@ -2,7 +2,7 @@ You are a general and precise AI assistant. I will ask you a question.
 Report your thoughts, and finish
 your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
 If a tool provide an error, use the tool differently.
-For web searching, ensure your answer by cross-checking data with several sources.
+For web searching, first search in your knowledge and if necessary complete them with web_search and ensure your answer by cross-checking data with several sources.
 YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of
 numbers and/or strings.
 If you are asked for a number, don’t use comma to write your number neither use units such as $ or percent

src/inference.py

@@ -19,6 +19,7 @@ from src.tools import (
     analyze_excel,
     analyze_youtube_video,
     calculate_sum,
+    retrieve_knowledge,
 )
 
 load_dotenv()
@@ -43,6 +44,7 @@ class Agent:
             "analyze_excel": analyze_excel,
             "analyze_youtube_video": analyze_youtube_video,
             "calculate_sum": calculate_sum,
+            "retrieve_knowledge": retrieve_knowledge,
         }
         self.log = []
         self.tools = self.get_tools()
@@ -74,6 +76,7 @@ class Agent:
                 analyze_excel,
                 analyze_youtube_video,
                 calculate_sum,
+                retrieve_knowledge,
             ]
         ).get('tools')
 

src/tools/__init__.py

@@ -9,4 +9,5 @@ from .transcribe_audio import transcribe_audio
 from .execute_code import execute_code
 from .analyze_excel import analyze_excel
 from .analyze_youtube_video import analyze_youtube_video
-from .calculator import calculate_sum
+from .calculator import calculate_sum
+from .retrieve_knowledge import retrieve_knowledge

src/tools/retrieve_knowledge.py

@@ -0,0 +1,22 @@
+from src.utils.tooling import tool
+
+@tool
+def retrieve_knowledge(query: str, n_results: int = 5, distance_threshold : float = 0.5) -> str:
+    """
+    Retrieves knowledge from a database with a provided query.
+    Args:
+        query (str): The query to search for in the vector store.
+        n_results (int, optional): The number of results to return. Default is 5.
+        distance_threshold (float, optional): The minimum distance score for results. Default is 0.5.
+    """
+    try:
+        from src.utils.vector_store import retrieve_from_database
+        results = retrieve_from_database(
+            query=query,
+            n_results=n_results,
+            distance_threshold=distance_threshold
+        )
+        return str(results)
+
+    except Exception as e:
+        return f"An unexpected error occurred: {str(e)}"

src/tools/visit_webpage.py

@@ -1,6 +1,5 @@
-import re
-
 from src.utils.tooling import tool
+from src.utils.vector_store import vectorize, load_in_vector_db
 
 @tool
 def visit_webpage(url: str) -> str:
@@ -11,6 +10,7 @@ def visit_webpage(url: str) -> str:
         url (str): The URL of the webpage to visit.
     """
     try:
+        import re
         import requests
         from markdownify import markdownify
         from requests.exceptions import RequestException
@@ -28,6 +28,14 @@ def visit_webpage(url: str) -> str:
         markdown_content = markdownify(response.text).strip()                            # Convert the HTML content to Markdown
         markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)      # Remove multiple line breaks
 
+        # Adding metadata
+        metadatas = {
+            "url": url,
+        }
+
+        text_embeddings, chunks = vectorize(markdown_content)   # Vectorize the content
+        load_in_vector_db(text_embeddings, chunks, metadatas=metadatas)              # Load the text embeddings into a FAISS index
+
         return truncate_content(markdown_content, 10000)
 
     except requests.exceptions.Timeout:
@@ -37,4 +45,4 @@ def visit_webpage(url: str) -> str:
         return f"Error fetching the webpage: {str(e)}"
 
     except Exception as e:
-        return f"An unexpected error occurred: {str(e)}"
+        return f"An unexpected error occurred: {str(e)}"

src/utils/vector_store.py

@@ -0,0 +1,124 @@
+import os
+from dotenv import load_dotenv
+from mistralai import Mistral
+import numpy as np
+import time
+import chromadb
+import json
+
+load_dotenv()
+MISTRAL_API_KEY = os.getenv("MISTRAL_API_KEY")
+COLLECTION_NAME = "webpages_collection"
+PERSIST_DIRECTORY = "./chroma_db"
+
+def get_text_embeddings(input_texts):
+    """
+    Get the text embeddings for the given inputs using Mistral API.
+    """
+    client = Mistral(api_key=MISTRAL_API_KEY)
+    while True:
+        try:
+            embeddings_batch_response = client.embeddings.create(
+                model="mistral-embed",
+                inputs=input_texts
+            )
+            return [data.embedding for data in embeddings_batch_response.data]
+        except Exception as e:
+            if "rate limit exceeded" in str(e).lower():
+                print("Rate limit exceeded. Retrying after 1 second...")
+                time.sleep(1)
+            else:
+                raise
+
+def vectorize(markdown_content, chunk_size=2048):
+    """
+    Vectorizes the given markdown content into chunks of specified size.
+    """
+    chunks = [markdown_content[i:i + chunk_size] for i in range(0, len(markdown_content), chunk_size)]
+    text_embeddings = get_text_embeddings(chunks)
+    return np.array(text_embeddings), chunks
+
+def load_in_vector_db(text_embeddings, chunks, metadatas=None, collection_name=COLLECTION_NAME):
+    """
+    Load the text embeddings into a ChromaDB collection for efficient similarity search.
+    """
+    client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
+
+    # Check if the collection exists, if not, create it
+    if collection_name not in [col.name for col in client.list_collections()]:
+        collection = client.create_collection(collection_name)
+    else:
+        collection = client.get_collection(collection_name)
+
+    for embedding, chunk in zip(text_embeddings, chunks):
+        collection.add(
+            embeddings=[embedding],
+            documents=[chunk],
+            metadatas=[metadatas],
+            ids=[str(hash(chunk))]
+        )
+
+
+def see_database(collection_name=COLLECTION_NAME):
+    """
+    Load the ChromaDB collection and text chunks.
+    """
+    client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
+
+    if collection_name not in [col.name for col in client.list_collections()]:
+        print("Collection not found. Please ensure it is created.")
+        return
+
+    collection = client.get_collection(collection_name)
+
+    items = collection.get()
+
+    print(f"Type of items: {type(items)}")
+    print(f"Items: {items}")
+
+    for item in items:
+        print(f"Type of item: {type(item)}")
+        print(f"Item: {item}")
+
+        if isinstance(item, dict):
+            print(f"ID: {item.get('ids')}")
+            print(f"Document: {item.get('document')}")
+            print(f"Metadata: {item.get('metadata')}")
+        else:
+            print("Item is not a dictionary")
+
+        print("---")
+
+def retrieve_from_database(query, collection_name=COLLECTION_NAME, n_results=5, distance_threshold=None):
+    """
+    Retrieve the most similar documents from the vector store based on the query.
+    """
+    client = chromadb.PersistentClient(path=PERSIST_DIRECTORY)
+    collection = client.get_collection(collection_name)
+    query_embeddings = get_text_embeddings([query])
+    raw_results = collection.query(
+        query_embeddings=query_embeddings,
+        n_results=n_results,
+        include=["documents", "metadatas", "distances"]
+    )
+    if distance_threshold is not None:
+        filtered_results = {
+            "ids": [],
+            "distances": [],
+            "metadatas": [],
+            "documents": []
+        }
+        for i, distance in enumerate(raw_results['distances'][0]):
+            if distance >= distance_threshold:
+                filtered_results['ids'].append(raw_results['ids'][0][i])
+                filtered_results['distances'].append(distance)
+                filtered_results['metadatas'].append(raw_results['metadatas'][0][i])
+                filtered_results['documents'].append(raw_results['documents'][0][i])
+        results = filtered_results
+
+        if len(results['documents']) == 0:
+            return "No relevant data found in knowledge database, have you visited webpages?"
+        else:
+            return json.dumps(results, indent=4)
+    else:
+        return json.dumps(raw_results, indent=4)

tools.json

@@ -253,5 +253,32 @@
                 ]
             }
         }
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "retrieve_knowledge",
+            "description": "Retrieves knowledge from a database with a provided query.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "query": {
+                        "type": "string",
+                        "description": "The query to search for in the vector store."
+                    },
+                    "n_results": {
+                        "type": "integer",
+                        "description": "The number of results to return. Default is 5."
+                    },
+                    "similarity_threshold": {
+                        "type": "number",
+                        "description": "The minimum similarity score for results. Default is 0.7."
+                    }
+                },
+                "required": [
+                    "query"
+                ]
+            }
+        }
     }
 ]