moved logic to engine

main.py

@@ -33,65 +33,23 @@ async def help(ctx: discord.ApplicationContext):
                       \n2)Those that tag the bot.")
 
 
-async def llm_output(question: str, context: str) -> str:
-    """
-    Returns output from the LLM using the given user-question and retrived context
-    """
-
-    URL_LLM = 'https://robinroy03-fury-bot.hf.space'
-    # URL_LLM = 'http://localhost:11434'    # NOTE: FOR TESTING
+async def llm_output(question: str) -> tuple[str, str]:
 
-    prompt = f"""
-    You are a senior FURY developer. FURY is a high level python graphics API similar to VTK.
-
-    Question: {question}
-
-    Context: {context}
-    """
+    URL_LLM = 'https://robinroy03-fury-engine.hf.space'
     obj = {
-            "model": "llama3-70b-8192",
-            "prompt": prompt,
-            "stream": False
-        }
+        "query": question,
+        "llm": "llama3-70b-8192",
+        "knn": 3,
+        "stream": False
+    }
 
     async with aiohttp.ClientSession() as session:
         async with session.post(URL_LLM + "/api/generate", json=obj) as response:
             if response.status == 500:
-                return "Sorry, an internal Error happened. Please try again later.\nError 500."
+                return "Error 500"
             response_json = await response.json()
 
-    return response_json['choices'][0]['message']['content']
-
-async def embedding_output(message: str) -> list:
-    """
-    Returns embeddings for the given message
-
-    rtype: list of embeddings. Length depends on the model.
-    """
-
-    URL_EMBEDDING = 'https://robinroy03-fury-embeddings-endpoint.hf.space'
-
-    async with aiohttp.ClientSession() as session:
-        async with session.post(URL_EMBEDDING + "/embedding", json={"text": message}) as response:
-            response_json = await response.json(content_type=None)
-    
-    return response_json['output']
-
-
-async def db_output(embedding: list) -> dict:
-    """
-    Returns the KNN results.
-
-    rtype: JSON
-    """
-
-    URL_DB = 'https://robinroy03-fury-db-endpoint.hf.space'
-    
-    async with aiohttp.ClientSession() as session:
-        async with session.post(URL_DB + "/query", json={"embeddings": embedding}) as response:
-            response_json = await response.json()
-        
-    return response_json
+    return response_json['response'], response_json['references']
 
 
 @bot.event
@@ -107,28 +65,9 @@ async def on_message(message):
     await message.reply(content="Your message was received, it'll take around 30 seconds for FURY to process an answer.")
 
     question = message.content.replace("<@1243428204124045385>", "")
-    
-    embedding: list = await embedding_output(question)
-
-    db_knn: dict = await db_output(embedding)
-    db_context = ""
-    references = ""
-    for i in range(len(db_knn['matches'])):
-        data = db_knn['matches'][i]['metadata']['data']
-        db_context += (data + "\n")
-        data = ast.literal_eval(data)
-        references += ("<https://github.com/fury-gl/fury/tree/master/" + data['path'] + ">").replace("//home/robin/Desktop/l/fury", "")
-        if data.get("function_name"):
-            references += f"\tFunction Name: {data.get('function_name')}"
-        elif data.get("class_name"):
-            references += f"\tClass Name: {data.get('class_name')}"
-        elif data['type'] == 'rst':
-            references += f"\tDocumentation: {data['path'].split("/")[-1]}"
-        elif data['type'] == 'documentation_examples':
-            references += f"\tDocumentation: {data['path'].split("/")[-1]}"
-        references += "\n"
-
-    llm_answer: str = await llm_output(question, db_context)
+    engine_response = await llm_output(question)
+    if llm_output != "Error 500":
+        llm_answer, references = engine_response
 
     try:
         start = 0
@@ -141,7 +80,7 @@ async def on_message(message):
         await message.reply(content=f"**References**\n{references}", view=Like_Dislike())                                  
     except Exception as e:                                                           # TODO: make exception handling better
         print(e)
-        await message.reply("An error occurred. Retry again.")
+        await message.reply("An error occurred. Retry again. Try different prompts.")
 
 def run_bot():
     bot.run(token)