modified to implement router

app.py

@@ -1,7 +1,7 @@
 """
 You call this ENDPOINT and it returns you a JSON which is of this format:
 
-POST FORMAT:
+POST FORMAT: (/api/groq or api/google or /api/ollama ...)
 {
     "query": "????",
     "llm": "llama70b-whatever",
@@ -12,88 +12,41 @@ POST FORMAT:
 RESPONSE FORMAT:
 {
     "response": "blabla",
-    "references": ["1", "2", "3", ...]
+    "references": "1, 2, 3"
 }
 """
 
 # TODO: MOVE IT ALL TO ASYNC FASTAPI, FOR NOW THIS IS A QUICK SPIN UP (IMPORTANT FOR SCALING)
 
-import ast
-import json
-
-import aiohttp
 from flask import Flask
 from flask import request
-import requests         # NOTE: AIOHTTP, THIS IS FOR RIGHT NOW ONLY
 
-app = Flask(__name__)
+from utils import embedding_output, db_output, groq_llm_output, ollama_llm_output
 
-def llm_output(question: str, db_knn: dict, llm: str, stream: bool) -> tuple[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
-
-    context = ""
-    references = ""
-    for i in range(len(db_knn['matches'])):
-        data = db_knn['matches'][i]['metadata']['data']
-        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"
-
-    prompt = f"""
-    You are a senior developer. Answer the users question based on the context provided.
-    Question: {question}
-    Context: {context}
-    """
-    obj = {
-            'model': llm,
-            'prompt': prompt,
-            'stream': stream
-        }
-    response = requests.post(URL_LLM + "/api/generate", json=obj)
-    response_json = json.loads(response.text)
-    return (response_json['choices'][0]['message']['content'], references)
 
+app = Flask(__name__)
 
-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'
-    response = requests.post(URL_EMBEDDING + "/embedding", json={"text": message})
-    response_json = json.loads(response.text)
-    return response_json['output']
+@app.route("/api/groq/generate", methods=['POST'])
+def groq_completion():
+    message = request.get_json()
 
+    query: str = message['query']
+    llm: str = message['llm']
+    knn: int = int(message['knn'])
+    stream: bool = bool(message['stream'])
 
-def db_output(embedding: list, knn: int) -> dict:
-    """
-    Returns the KNN results.
-    rtype: JSON
-    """
+    embedding_data = embedding_output(query)
+    db_knn = db_output(embedding_data, knn)
+    output, references = groq_llm_output(query, db_knn, llm, stream)
 
-    URL_DB = 'https://robinroy03-fury-db-endpoint.hf.space'
-    response = requests.post(URL_DB + "/query", json={"embeddings": embedding, "knn": knn})
-    response_json = json.loads(response.text)
-    return response_json
+    return {
+        "response": output,
+        "references": references
+    }
 
 
-@app.route("/api/generate", methods=['POST'])
-def completion():
+@app.route("/api/ollama/generate", methods=['POST'])
+def ollama_completion():
     message = request.get_json()
 
     query: str = message['query']
@@ -103,18 +56,24 @@ def completion():
 
     embedding_data = embedding_output(query)
     db_knn = db_output(embedding_data, knn)
-    output, references = llm_output(query, db_knn, llm, stream)
+    response_json, references = ollama_llm_output(query, db_knn, llm, stream)
+
+    if response_json.get("error"):
+        print(response_json)
+        return {
+            "response": "An error occured, try again.",
+            "references": "No references"
+        }
 
     return {
-        "response": output,
+        "response": response_json['response'],
         "references": references
     }
 
-
 """
-curl -X POST http://localhost:8000/api/generate -H "Content-Type: application/json" -d '{
+curl -X POST http://localhost:8000/api/ollama/generate -H "Content-Type: application/json" -d '{
     "query": "How do I create a sphere in FURY?",
-    "llm": "llama3-70b-8192",
+    "llm": "phi3",
     "knn": "3",
     "stream": false
   }'

utils.py

@@ -0,0 +1,89 @@
+import requests
+import json
+import ast
+
+
+def prompt_generator(question: str, db_knn: dict) -> tuple[str, str, str]:
+    context = ""
+    references = ""
+    for i in range(len(db_knn['matches'])):
+        data = db_knn['matches'][i]['metadata']['data']
+        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"
+
+    prompt = f"""
+    You are a senior developer. Answer the users question based on the context provided.
+    Question: {question}
+    Context: {context}
+    """
+    return prompt, context, references
+
+
+def groq_llm_output(question: str, db_knn: dict, llm: str, stream: bool) -> tuple[str, str]:
+    """
+    Returns output from the LLM using the given user-question and retrived context
+    """
+
+    URL_LLM = 'https://robinroy03-fury-bot.hf.space'
+    prompt, context, references = prompt_generator(question, db_knn)
+    obj = {
+            'model': llm,
+            'prompt': prompt,
+            'stream': stream
+        }
+    response = requests.post(URL_LLM + "/api/generate", json=obj)
+    response_json = json.loads(response.text)
+    return (response_json['choices'][0]['message']['content'], references)
+
+
+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'
+    response = requests.post(URL_EMBEDDING + "/embedding", json={"text": message})
+    response_json = json.loads(response.text)
+    return response_json['output']
+
+
+def db_output(embedding: list, knn: int) -> dict:
+    """
+    Returns the KNN results.
+    rtype: JSON
+    """
+
+    URL_DB = 'https://robinroy03-fury-db-endpoint.hf.space'
+    response = requests.post(URL_DB + "/query", json={"embeddings": embedding, "knn": knn})
+    response_json = json.loads(response.text)
+    return response_json
+
+
+def ollama_llm_output(question: str, db_knn: dict, llm: str, stream: bool) -> tuple[str, str]:
+    # URL_LLM = 'https://robinroy03-ollama-server-backend.hf.space'
+    URL_LLM = "http://localhost:11434"
+    prompt, context, references = prompt_generator(question, db_knn)
+    obj = {
+        "model": llm,
+        "prompt": question,
+        "stream": stream
+    }
+    try:
+        response = requests.post(URL_LLM + "/api/generate", json=obj)
+    except Exception as e:
+        print(e)
+        return {"error": e}
+
+    response_json = json.loads(response.text)
+    return response_json, references