added websocket to hf_mixtral_agent

innovation_pathfinder_ai/backend/app/api/v1/agents/hf_mixtral_agent.py

@@ -7,14 +7,14 @@ from langchain.agents.output_parsers import ReActJsonSingleInputOutputParser
 from langchain.tools.render import render_text_description
 import os
 from dotenv import load_dotenv
-from innovation_pathfinder_ai.backend.app.structured_tools.structured_tools import (
+from app.structured_tools.structured_tools import (
     arxiv_search, get_arxiv_paper, google_search, wikipedia_search, knowledgeBase_search, memory_search
 )
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect
 from langchain.prompts import PromptTemplate
-from innovation_pathfinder_ai.backend.app.templates.react_json_with_memory import template_system
-from innovation_pathfinder_ai.backend.app.utils import logger
-from innovation_pathfinder_ai.backend.app.utils import generate_uuid
+from app.templates.react_json_with_memory import template_system
+from app.utils import logger
+from app.utils import utils
 from langchain.globals import set_llm_cache
 from langchain.cache import SQLiteCache
 
@@ -32,7 +32,7 @@ LANGCHAIN_PROJECT = os.getenv('LANGCHAIN_PROJECT')
 
 router = APIRouter()
 
-@router.websocket("")
+@router.websocket("/agent")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
 
@@ -41,16 +41,16 @@ async def websocket_endpoint(websocket: WebSocket):
             data = await websocket.receive_json()
             user_message = data["message"]
 
-            resp = IChatResponse(
-                sender="you",
-                message=user_message_card.to_dict(),
-                type="start",
-                message_id=generate_uuid(),
-                id=generate_uuid(),
-            )
+            # resp = IChatResponse(
+            #     sender="you",
+            #     message=user_message_card.to_dict(),
+            #     type="start",
+            #     message_id=generate_uuid(),
+            #     id=generate_uuid(),
+            # )
 
-            await websocket.send_json(resp.dict())
-            message_id: str = generate_uuid()
+            # await websocket.send_json(resp.dict())
+            message_id: str = utils.generate_uuid()
             # custom_handler = CustomFinalStreamingStdOutCallbackHandler(
             #     websocket, message_id=message_id
             # )

innovation_pathfinder_ai/backend/app/api/v1/endpoints/add_to_kb.py

@@ -1,8 +1,8 @@
 from fastapi import APIRouter
-from innovation_pathfinder_ai.backend.app.utils.utils import extract_urls
-from innovation_pathfinder_ai.backend.app.utils import logger
-from innovation_pathfinder_ai.backend.app.vector_store import initialize_chroma_db 
-from innovation_pathfinder_ai.backend.app.utils.utils import (
+from app.utils.utils import extract_urls
+from app.utils import logger
+from app.vector_store import initialize_chroma_db 
+from app.utils.utils import (
     generate_uuid    
 )
 from langchain_community.vectorstores import Chroma

innovation_pathfinder_ai/backend/app/crud/db_handler.py

@@ -1,6 +1,6 @@
 from sqlmodel import SQLModel, create_engine, Session, select
-from innovation_pathfinder_ai.backend.app.database.db_schema import Sources
-from innovation_pathfinder_ai.backend.app.utils.logger import get_console_logger
+from app.database.db_schema import Sources
+from app.utils.logger import get_console_logger
 import os
 from dotenv import load_dotenv
 

innovation_pathfinder_ai/backend/app/main.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI
-from innovation_pathfinder_ai.backend.app.api.v1.api import api_router as api_router_v1
+from app.api.v1.api import api_router as api_router_v1
 #from app.core.config import settings
 from fastapi.middleware.cors import CORSMiddleware
 

innovation_pathfinder_ai/backend/app/structured_tools/structured_tools.py

@@ -15,24 +15,24 @@ import ast
 import chromadb
 
 # hacky and should be replaced with a database
-from innovation_pathfinder_ai.source_container.container import (
-    all_sources
-)
-from innovation_pathfinder_ai.utils.utils import (
+# from app.source_container.container import (
+#     all_sources
+# )
+from app.utils.utils import (
     parse_list_to_dicts, format_wiki_summaries, format_arxiv_documents, format_search_results
 )
-from backend.app.crud.db_handler import (
+from app.crud.db_handler import (
     add_many
 )
 
-from innovation_pathfinder_ai.vector_store.chroma_vector_store import (
+from app.vector_store.chroma_vector_store import (
     add_pdf_to_vector_store
 )
-from innovation_pathfinder_ai.utils.utils import (
+from app.utils.utils import (
     create_wikipedia_urls_from_text, create_folder_if_not_exists,
 )
 import os
-# from innovation_pathfinder_ai.utils import create_wikipedia_urls_from_text
+# from app.utils import create_wikipedia_urls_from_text
 
 persist_directory = os.getenv('VECTOR_DATABASE_LOCATION')
 
@@ -93,14 +93,14 @@ def knowledgeBase_search(query:str) -> str:
 def arxiv_search(query: str) -> str:
     """Search arxiv database for scientific research papers and studies. This is your primary online information source.
     always check it first when you search for additional information, before using any other online tool."""
-    global all_sources
+    #global all_sources
     arxiv_retriever = ArxivRetriever(load_max_docs=3)
     data = arxiv_retriever.invoke(query)
     meta_data = [i.metadata for i in data]
     formatted_sources = format_arxiv_documents(data)
-    all_sources += formatted_sources
+    #all_sources += formatted_sources
     parsed_sources = parse_list_to_dicts(formatted_sources)
-    add_many(parsed_sources)
+    #add_many(parsed_sources)
   
     return data.__str__()
 
@@ -162,28 +162,28 @@ def embed_arvix_paper(paper_id:str) -> None:
 @tool
 def wikipedia_search(query: str) -> str:
     """Search Wikipedia for additional information to expand on research papers or when no papers can be found."""
-    global all_sources
+    #global all_sources
 
     api_wrapper = WikipediaAPIWrapper()
     wikipedia_search = WikipediaQueryRun(api_wrapper=api_wrapper)
     wikipedia_results = wikipedia_search.run(query)
     formatted_summaries = format_wiki_summaries(wikipedia_results)
-    all_sources += formatted_summaries
+    #all_sources += formatted_summaries
     parsed_summaries = parse_list_to_dicts(formatted_summaries)
-    add_many(parsed_summaries)
+    #add_many(parsed_summaries)
     #all_sources += create_wikipedia_urls_from_text(wikipedia_results)
     return wikipedia_results
 
 @tool
 def google_search(query: str) -> str:
     """Search Google for additional results when you can't answer questions using arxiv search or wikipedia search."""
-    global all_sources
+    #global all_sources
     
     websearch = GoogleSearchAPIWrapper()
     search_results:dict = websearch.results(query, 3)
     cleaner_sources =format_search_results(search_results)
     parsed_csources = parse_list_to_dicts(cleaner_sources)
-    add_many(parsed_csources)
-    all_sources += cleaner_sources    
+    #add_many(parsed_csources)
+    #all_sources += cleaner_sources    
     
     return cleaner_sources.__str__()

innovation_pathfinder_ai/backend/app/utils/utils.py

@@ -3,7 +3,7 @@ import datetime
 import os
 import uuid
 
-from innovation_pathfinder_ai.backend.app.utils import logger
+from app.utils import logger
 
 logger = logger.get_console_logger("utils")
 

innovation_pathfinder_ai/backend/app/vector_store/chroma_vector_store.py

@@ -20,7 +20,7 @@ from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings.sentence_transformer import (
     SentenceTransformerEmbeddings,
 )
-from innovation_pathfinder_ai.utils.utils import (
+from app.utils.utils import (
     generate_uuid    
 )
 import dotenv

innovation_pathfinder_ai/frontend/app.py

@@ -1,47 +1,34 @@
 from fastapi import FastAPI
 import gradio as gr
 from gradio.themes.base import Base
-from innovation_pathfinder_ai.backend.app.api.v1.agents.hf_mixtral_agent import agent_executor
-from innovation_pathfinder_ai.source_container.container import (
-    all_sources
-)
-from innovation_pathfinder_ai.backend.app.utils.utils import extract_urls
-from innovation_pathfinder_ai.backend.app.utils import logger
-
-from innovation_pathfinder_ai.backend.app.utils.utils import (
-    generate_uuid    
-)
+#from innovation_pathfinder_ai.backend.app.api.v1.agents.hf_mixtral_agent import agent_executor
+#from innovation_pathfinder_ai.source_container.container import (
+#    all_sources
+#)
+#from innovation_pathfinder_ai.backend.app.utils.utils import extract_urls
+#from innovation_pathfinder_ai.backend.app.utils import logger
+#from innovation_pathfinder_ai.backend.app.vector_store.chroma_vector_store import initialize_chroma_db
+#from innovation_pathfinder_ai.backend.app.utils.utils import (
+#    generate_uuid    
+#)
 from langchain_community.vectorstores import Chroma
 
-import chromadb
+import asyncio
+import websockets
+import json
 import dotenv
 import os
 
 dotenv.load_dotenv()
 persist_directory = os.getenv('VECTOR_DATABASE_LOCATION')
 
-logger = logger.get_console_logger("app")
+#logger = logger.get_console_logger("app")
 
 app = FastAPI()
 
-def initialize_chroma_db() -> Chroma:
-    collection_name = os.getenv('CONVERSATION_COLLECTION_NAME')
-    
-    client = chromadb.PersistentClient(
-        path=persist_directory
-        )
-    
-    collection = client.get_or_create_collection(
-    name=collection_name,
-    )
-    
-    return collection
-
-
-
 if __name__ == "__main__":
     
-    db = initialize_chroma_db()
+    #db = initialize_chroma_db()
     
     def add_text(history, text):
         history = history + [(text, None)]
@@ -53,35 +40,51 @@ if __name__ == "__main__":
         # Example for calling generate_uuid from the backend
         # response = requests.post("http://localhost:8000/add-document")
         #current_id = response.text
-        sources = extract_urls(all_sources)
-        src_list = '\n'.join(sources)
-        current_id = generate_uuid()
-        db.add(
-            ids=[current_id],
-            documents=[response['output']],
-            metadatas=[
-                {
-                    "human_message":history[-1][0],
-                    "sources": 'Internal Knowledge Base From: \n\n' + src_list
-                }
-            ]
-        )
-        if not sources:
-            response_w_sources = response['output']+"\n\n\n Sources:  \n\n\n Internal knowledge base"
-        else:
-            response_w_sources = response['output']+"\n\n\n Sources: \n\n\n"+src_list
-        history[-1][1] = response_w_sources
-        all_sources.clear()
+        # sources = extract_urls(all_sources)
+        # src_list = '\n'.join(sources)
+        # current_id = generate_uuid()
+        # db.add(
+        #     ids=[current_id],
+        #     documents=[response['output']],
+        #     metadatas=[
+        #         {
+        #             "human_message":history[-1][0],
+        #             "sources": 'Internal Knowledge Base From: \n\n' + src_list
+        #         }
+        #     ]
+        # )
+        # if not sources:
+        #     response_w_sources = response['output']+"\n\n\n Sources:  \n\n\n Internal knowledge base"
+        # else:
+        #     response_w_sources = response['output']+"\n\n\n Sources: \n\n\n"+src_list
+        history[-1][1] = response['output']
+       # all_sources.clear()
         return history
 
     def infer(question, history):
-        query =  question
-        result = agent_executor.invoke(
-            {
-                "input": question,
-                "chat_history": history
-            }
-        )        
+        # result = agent_executor.invoke(
+        #     {
+        #         "input": question,
+        #         "chat_history": history
+        #     }
+        # )        
+        # return result
+        async def ask_question_async(question, history):
+            uri = "ws://localhost:8000/chat/agent"  # Update this URI to your actual WebSocket endpoint
+            async with websockets.connect(uri) as websocket:
+                # Prepare the message to send (adjust the structure as needed for your backend)
+                message_data = {
+                    "message": question,
+                    "history": history
+                }
+                await websocket.send(json.dumps(message_data))
+
+                # Wait for the response
+                response_data = await websocket.recv()
+                return json.loads(response_data)
+
+        # Run the asynchronous function in the synchronous context
+        result = asyncio.get_event_loop().run_until_complete(ask_question_async(question, history))
         return result
 
     def vote(data: gr.LikeData):
@@ -122,7 +125,7 @@ if __name__ == "__main__":
                 gr.Markdown("Nothing yet...")
 
     demo.queue()
-    demo.launch(debug=True, favicon_path="innovation_pathfinder_ai/assets/favicon.ico", share=True)
+    demo.launch(debug=True, favicon_path="assets/favicon.ico", share=True)
 
     x = 0 # for debugging purposes
     app = gr.mount_gradio_app(app, demo, path="/")