new google gemini added

app.py

@@ -21,7 +21,7 @@ RESPONSE FORMAT:
 from flask import Flask
 from flask import request
 
-from utils import embedding_output, db_output, groq_llm_output, ollama_llm_output
+from utils import embedding_output, db_output, groq_llm_output, ollama_llm_output, google_llm_output
 
 
 app = Flask(__name__)
@@ -70,6 +70,26 @@ def ollama_completion():
         "references": references
     }
 
+
+@app.route("/api/google/generate", methods=['POST'])
+def google_completion():
+    message = request.get_json()
+
+    query: str = message['query']
+    llm: str = message['llm']
+    knn: int = int(message['knn'])
+    stream: bool = bool(message['stream'])
+
+    embedding_data = embedding_output(query)
+    db_knn = db_output(embedding_data, knn)
+    response_json, references = google_llm_output(query, db_knn, llm, stream)
+
+    return {
+        "response": response_json['response'],
+        "references": references
+    }
+
+
 """
 curl -X POST http://localhost:8000/api/groq/generate -H "Content-Type: application/json" -d '{
     "query": "How do I create a sphere in FURY?",

utils.py

@@ -41,11 +41,24 @@ def groq_llm_output(question: str, db_knn: dict, llm: str, stream: bool) -> tupl
             'prompt': prompt,
             'stream': stream
         }
-    response = requests.post(URL_LLM + "/api/generate", json=obj)
+    response = requests.post(URL_LLM + "/api/groq/generate", json=obj)
     response_json = json.loads(response.text)
     return (response_json['choices'][0]['message']['content'], references)
 
 
+def google_llm_output(question: str, db_knn: dict, llm: str, stream: bool) -> tuple[str, str]:
+    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/google/generate", json=obj)
+    response_json = json.loads(response.text)
+    return (response_json['candidates'][0]['content']['parts'][0]['text'], references)
+
+
 def embedding_output(message: str) -> list:
     """
     Returns embeddings for the given message