Update data_collector.py

data_collector.py

@@ -6,22 +6,63 @@ from sentence_transformers import SentenceTransformer,util
 from transformers import pipeline
 import requests
 
-def consume_llm_api(prompt):
-    """
-    Sends a prompt to the LLM API and processes the streamed response.
-    """
-    url = "https://8417-201-238-124-65.ngrok-free.app/api/llm-response"
-    headers = {"Content-Type": "application/json"}
-    payload = {"prompt": prompt,"extension":"1"}
+# def consume_llm_api(prompt):
+#     """
+#     Sends a prompt to the LLM API and processes the streamed response.
+#     """
+#     url = "https://8417-201-238-124-65.ngrok-free.app/api/llm-response"
+#     headers = {"Content-Type": "application/json"}
+#     payload = {"prompt": prompt,"extension":"1"}
 
     
-    print("Sending prompt to the LLM API...")
-    response_ = requests.post(url, json=payload,verify=False)
-    response_data = response_.json()
-    return response_data['text']
-# def consume_llm_api(prompt):
-#     model = Ollama(model="llama3:latest", temperature=0.3)
-#     return model.invoke(prompt)
+#     print("Sending prompt to the LLM API...")
+#     response_ = requests.post(url, json=payload,verify=False)
+#     response_data = response_.json()
+#     return response_data['text']
+def consume_llm_api(prompt):
+
+    import requests
+    from sentence_transformers import SentenceTransformer
+    from pinecone import Pinecone, ServerlessSpec
+    Gen_api = "https://8417-201-238-124-65.ngrok-free.app/api/llm-response"
+    api_key = "pcsk_2EhvKP_GqGkpAjF4p4ziL7PgrgM9xuKcthX9gtqhyLxV3UaMmWTQufW4qKZrjhLrf2d1ma"
+    pc = Pinecone(api_key=api_key)
+    model = SentenceTransformer("all-mpnet-base-v2")
+    try:
+        index_name = "quickstart"
+        pc.create_index(
+            name=index_name,
+            dimension=768, 
+            metric="cosine", 
+            spec=ServerlessSpec(
+                cloud="aws",
+                region="us-east-1"
+            ) 
+        )
+    except:
+        pass
+    index = pc.Index(index_name)
+    index.upsert(
+        vectors=[
+            {
+                "id": "lorum", 
+                "values": [float(i) for i in list(model.encode("lorum"))],  
+                "metadata": {"string":str(prompt)}
+                
+            }
+        ]
+    )
+
+    gen_api_response = requests.post(url = Gen_api,json={"api_key": api_key},verify=False)
+
+    if gen_api_response.json().get("status"):
+        response = index.query(
+        vector=[float(i) for i in model.encode(str(prompt))],
+        top_k=1,
+        include_metadata=True,
+    )
+        
+    return response['matches'][0]['metadata']['string']
 
 
 def relevent_value(long_query,count=3):