Update gemini_fastapi_server.py

gemini_fastapi_server.py

@@ -1,7 +1,8 @@
 import os
 from fastapi import FastAPI
 from langchain.prompts import ChatPromptTemplate
-from langchain_gemini import ChatGemini  # Make sure this is the correct import for Gemini
+from langchain_gemini import ChatGemini  # Assuming there’s a Gemini integration
+from langchain import HuggingFaceChat  # Hugging Face integration
 from langserve import add_routes
 import uvicorn
 import nest_asyncio
@@ -9,28 +10,41 @@ import nest_asyncio
 # Allows the use of uvicorn.run in environments like Jupyter
 nest_asyncio.apply()
 
+  # Retrieve HF token
+gemini_api_key = os.getenv("AIzaSyBo5SPd4H4gM0ONHBXRHAYoc973szdrfk4")  # Retrieve Gemini API key
+
 # Initialize the FastAPI app
 app = FastAPI(
-    title="Langchain Server",
+    title="Multimodal Language Server",
     version="1.0",
-    description="A simple QnA API Server using LLaMA 3.2 from Hugging Face"
+    description="A simple QnA API Server using both Hugging Face and Gemini models"
 )
 
-# Initialize the LLaMA model using the Hugging Face chat model
-hf_token = os.environ['HF_TOKEN']  # Get the HF API token
-llama_model = HuggingFaceChat(model="meta-llama/LLaMA-3-2", token=hf_token)  # Use the correct model identifier
+# Initialize the LLaMA model using Hugging Face
+llama_model = HuggingFaceChat(model="meta-llama/LLaMA-3-2", token=HF_TOKEN)  # Use the specific LLaMA model from HF
+
+# Initialize the Gemini model (adjust based on the actual integration)
+gemini_model = ChatGemini(api_key=gemini_api_key, model="gemini_model_name_here")  # Specify the correct model name
 
 # Define a QnA prompt using a template
 qna_prompt = ChatPromptTemplate.from_template("Answer the question: {question}")
 
-# Add a route for the QnA API using the LLaMA model
+# Function to choose model based on preference
+# You might want to improve this to include more refined selection criteria
+def get_model_response(question, use_gemini=False):
+    if use_gemini:
+        return gemini_model(question)  # Call the Gemini model
+    else:
+        return llama_model(question)  # Call the Hugging Face model
+
+# Create an API endpoint
 @app.post("/llm_api")
-async def qna_endpoint(question: str):
-add_routes(
-    app,
-    qna_prompt | llama_model,
-    path="/llm_api"
-)
+async def qna_endpoint(question: str, use_gemini: bool = False):
+    """
+    Endpoint to receive a question and return a response from either the Hugging Face or Gemini model.
+    """
+    response = get_model_response(question, use_gemini)
+    return {"response": response}
 
 # Run the application
 if __name__ == "__main__":