Update app.py

app.py

@@ -11,7 +11,7 @@ from apscheduler.schedulers.background import BackgroundScheduler
 import os
 from pydantic import BaseModel
 
-SAMBA_NOVA_API_KEY = os.environ.get("SAMBA_NOVA_API_KEY", None)
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", None)
 
 app = FastAPI()
 
@@ -34,7 +34,7 @@ scheduler.start()
 class StreamTextRequest(BaseModel):
     query: str
     history: str = "[]"
-    model: str = "llama3-8b"
+    model: str = "gpt-4"
     api_key: str = None
 
 @app.post("/stream_text")
@@ -48,29 +48,25 @@ async def stream_text(request: StreamTextRequest):
         return StreamingResponse(iter([f"{cached_response}"]), media_type='text/event-stream')
 
     # Model selection logic
-    if "405" in request.model:
-        fmodel = "Meta-Llama-3.1-405B-Instruct"
-    if "70" in request.model:
-        fmodel = "Meta-Llama-3.1-70B-Instruct"
-    else:
-        fmodel = "Meta-Llama-3.1-8B-Instruct"
+    fmodel = "gpt-4" if "gpt-4" in request.model else "gpt-3.5-turbo"
 
     system_message = """You are an AI language assistant. Your sole task is to correct the grammar, spelling, and structure of the sentences provided to you. You must not change the meaning of the sentences, and you should focus only on making them grammatically correct, concise, and clear. Do not add any additional information or provide explanations unless specifically asked. Your responses should be limited to the corrected version of the sentence."""
 
     messages = [{'role': 'system', 'content': system_message}]
-
     messages.extend(ast.literal_eval(request.history))
-
     messages.append({'role': 'user', 'content': request.query})
     
-    data = {'messages': messages, 'stream': True, 'model': fmodel}
-
-    api_key = request.api_key if request.api_key != 'none' else SAMBA_NOVA_API_KEY
+    data = {'messages': messages, 'model': fmodel, 'stream': True}
 
+    api_key = request.api_key if request.api_key != 'none' else OPENAI_API_KEY
 
     async def stream_response():
         async with aiohttp.ClientSession() as session:
-            async with session.post('https://api.sambanova.ai/v1/chat/completions', headers = { 'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json' }, json=data) as response:
+            async with session.post(
+                'https://api.openai.com/v1/chat/completions',
+                headers={'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'},
+                json=data
+            ) as response:
                 if response.status != 200:
                     raise HTTPException(status_code=response.status, detail="Error fetching AI response")
 
@@ -95,10 +91,8 @@ async def stream_text(request: StreamTextRequest):
 
     return StreamingResponse(stream_response(), media_type='text/event-stream')
 
-
-
-# Serve index.html from the same directory as your main.py file 
-from starlette.responses import FileResponse 
+# Serve static files
+from starlette.responses import FileResponse
 
 @app.get("/script1.js")
 async def script1_js():
@@ -118,4 +112,4 @@ async def read_index():
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7068, reload=True)
+    uvicorn.run(app, host="0.0.0.0", port=7068, reload=True)