Updated rag-tester and web crawler

Author: nullchimp

examples/8-url-loader.py

@@ -11,13 +11,15 @@
 import os
 
 db = MemGraphClient(
-    host=os.environ.get("MEMGRAPH_URI", "localhost"),
+    host="localhost" or os.environ.get("MEMGRAPH_URI", "localhost"),
     port=int(os.environ.get("MEMGRAPH_PORT", 7687)),
     username=os.environ.get("MEMGRAPH_USERNAME", "memgraph"),
     password=os.environ.get("MEMGRAPH_PASSWORD", "memgraph"),
 )
 db.connect()
 
+print("Connected to Memgraph", db.host, db.port)
+
 loader = WebLoader("http://localhost:4000/en/enterprise-cloud@latest/copilot")
 embedder = TextEmbedding3Small()
 
@@ -26,6 +28,12 @@
 )
 
 def store(source, doc, chunks, vectors):
+    print("### Storing data in Memgraph")
+    print("Source:", source)
+    print("Document:", doc)
+    print("Chunks:", len(chunks))
+    print("Vectors:", len(vectors))
+
     db.create_source(source)
     db.create_document(doc)
     for chunk in chunks:
@@ -35,6 +43,8 @@ def store(source, doc, chunks, vectors):
         vector.vector_store_id = vector_store.id
         db.create_vector(vector)
 
+    print("### Data stored successfully")
+
 async def main():
     for source, doc, chunks in loader.load_data():
         vectors = []

examples/9-test.py

@@ -21,7 +21,7 @@
 
 # Define enhanced system role with instructions on using all available tools
 system_role = f"""
-You are an expert on everything GitHub.
+You are a helpful assistant. 
 Your Name is Agent Smith.
 
 Today is {date.today().strftime("%d %B %Y")}.
@@ -53,11 +53,11 @@ async def run_conversation(user_prompt: str, rag_prompt = None) -> str:
     raise ValueError(f"AZURE_OPENAI_API_KEY environment variable is required")
 
 db = MemGraphClient(
-    host=os.environ.get("MEMGRAPH_URI", "localhost"),
+    host="localhost" or os.environ.get("MEMGRAPH_URI", "localhost"),
     port=int(os.environ.get("MEMGRAPH_PORT", 7687)),
     username=os.environ.get("MEMGRAPH_USERNAME", "memgraph"),
     password=os.environ.get("MEMGRAPH_PASSWORD", "memgraph"),
-)
+).connect()
 
 embedder = TextEmbedding3Small()
 
@@ -129,15 +129,26 @@ async def test_vector_search(query_text: str):
         "references": <list of references to other sources>
     }}
 ]
-I need you to always ground your response in this information and return relevant sources and references.
+I always HAVE TO ground your response in this information.
+You always HAVE TO return the source, where the information is coming from, at the end of your response.
+Use the following format:
+
+# Sources:
+1. <Source Name>: <Source URI>
+2. <Source Name>: <Source URI>
+...
+
+It is totally fine to only use 1 source, but you have to mention it.
+If you don't know the answer, say "I don't know".
 Here is the information You have:\n
     """
 
     await run_conversation(query_text, f"{rag_prompt}{json.dumps(data)}")
 
 
 async def main():
-    text = input("Enter your text: ")
+    #text = input("Enter your text: ")
+    text = "Tell me what the multipliers for premium requests are, based on the Model for GitHub Copilot!"
     await test_vector_search(text)
     #await run_conversation(text)
     print("Test completed.")

src/core/rag/dbhandler/memgraph.py

@@ -42,4 +42,6 @@ def connect(self, *args, **kwargs) -> None:
         self._conn.autocommit = True
         self._cur = self._conn.cursor()
         print(f"Connected successfully to Memgraph")
+
+        return self
 

src/libs/dataloader/web.py

@@ -116,6 +116,7 @@ def _visit_site(self, url: str, retry = 3) -> str:
         except Exception as e:
             if retry <= 0:
                 raise ValueError(f"Failed to fetch content from {url}: {str(e)}")
+            print(f"Error fetching {url}: {str(e)}, Retrying... ({retry - 1} attempts left)")
             return self._visit_site(url, retry - 1)
 
     def load_data(self) -> Generator[Tuple[Source, Document, List[DocumentChunk]], None, None]: