Update app.py

app.py

@@ -4,100 +4,73 @@ from fastapi.responses import StreamingResponse
 from openai import AsyncOpenAI
 
 app = FastAPI()
+
+# System prompt
 system = '''You are DeepSeek R1, an advanced reasoning assistant.
 Your responses consist of two parts:
-
 1. A <thinking> block — This is your internal reasoning. You think step-by-step, carefully analyzing the question, considering context, alternatives, and edge cases. This section must be at least 10 lines long and enclosed between <think> and </think>. This part is not shown to the user in real-world applications, but is visible during debugging or development.
 2. The final answer — This is the polished, professional response provided after you’ve thought through the problem. It is clear, structured, and concise.
 3. always provide code in this foramte ```<code>```.
+
 Your behavior guidelines:
-Maintain a calm, analytical, and formal tone.
-Use bullet points or numbered lists when appropriate.
-Avoid casual language, emojis, or redundant filler.
-If context is missing, mention assumptions.
-Never refer to yourself as an AI or language model.
-Do not repeat the <thinking> part in your final answer.
+- Maintain a calm, analytical, and formal tone.
+- Use bullet points or numbered lists when appropriate.
+- Avoid casual language, emojis, or redundant filler.
+- If context is missing, mention assumptions.
+- Never refer to yourself as an AI or language model.
+- Do not repeat the <thinking> part in your final answer.
+
 Format every response exactly as follows:
 <think>  
 [Begin detailed, line-by-line reasoning here — minimum 10 lines. Think aloud.]  
 </think>  
-[Final answer starts here — no label, just a clean professional response.]  '''
-# Define available models (you can expand this list)
+[Final answer starts here — no label, just a clean professional response.]  
+'''
+
+# In-memory chat history
+chat_history = {}
+
+# Supported models
 AVAILABLE_MODELS = {
-  "openai/gpt-4.1": "OpenAI GPT-4.1",
-  "openai/gpt-4.1-mini": "OpenAI GPT-4.1-mini",
-  "openai/gpt-4.1-nano": "OpenAI GPT-4.1-nano",
-  "openai/gpt-4o": "OpenAI GPT-4o",
-  "openai/gpt-4o-mini": "OpenAI GPT-4o mini",
-  "openai/o4-mini": "OpenAI o4-mini",
-  "microsoft/MAI-DS-R1": "MAI-DS-R1",
-  "microsoft/Phi-3.5-MoE-instruct": "Phi-3.5-MoE instruct (128k)",
-  "microsoft/Phi-3.5-mini-instruct": "Phi-3.5-mini instruct (128k)",
-  "microsoft/Phi-3.5-vision-instruct": "Phi-3.5-vision instruct (128k)",
-  "microsoft/Phi-3-medium-128k-instruct": "Phi-3-medium instruct (128k)",
-  "microsoft/Phi-3-medium-4k-instruct": "Phi-3-medium instruct (4k)",
-  "microsoft/Phi-3-mini-128k-instruct": "Phi-3-mini instruct (128k)",
-  "microsoft/Phi-3-small-128k-instruct": "Phi-3-small instruct (128k)",
-  "microsoft/Phi-3-small-8k-instruct": "Phi-3-small instruct (8k)",
-  "microsoft/Phi-4": "Phi-4",
-  "microsoft/Phi-4-mini-instruct": "Phi-4-mini-instruct",
-  "microsoft/Phi-4-multimodal-instruct": "Phi-4-multimodal-instruct",
-  "ai21-labs/AI21-Jamba-1.5-Large": "AI21 Jamba 1.5 Large",
-  "ai21-labs/AI21-Jamba-1.5-Mini": "AI21 Jamba 1.5 Mini",
-  "mistral-ai/Codestral-2501": "Codestral 25.01",
-  "cohere/Cohere-command-r": "Cohere Command R",
-  "cohere/Cohere-command-r-08-2024": "Cohere Command R 08-2024",
-  "cohere/Cohere-command-r-plus": "Cohere Command R+",
-  "cohere/Cohere-command-r-plus-08-2024": "Cohere Command R+ 08-2024",
-  "deepseek/DeepSeek-R1": "DeepSeek-R1",
-  "deepseek/DeepSeek-V3-0324": "DeepSeek-V3-0324",
-  "meta/Llama-3.2-11B-Vision-Instruct": "Llama-3.2-11B-Vision-Instruct",
-  "meta/Llama-3.2-90B-Vision-Instruct": "Llama-3.2-90B-Vision-Instruct",
-  "meta/Llama-3.3-70B-Instruct": "Llama-3.3-70B-Instruct",
-  "meta/Llama-4-Maverick-17B-128E-Instruct-FP8": "Llama 4 Maverick 17B 128E Instruct FP8",
-  "meta/Llama-4-Scout-17B-16E-Instruct": "Llama 4 Scout 17B 16E Instruct",
-  "meta/Meta-Llama-3.1-405B-Instruct": "Meta-Llama-3.1-405B-Instruct",
-  "meta/Meta-Llama-3.1-70B-Instruct": "Meta-Llama-3.1-70B-Instruct",
-  "meta/Meta-Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct",
-  "meta/Meta-Llama-3-70B-Instruct": "Meta-Llama-3-70B-Instruct",
-  "meta/Meta-Llama-3-8B-Instruct": "Meta-Llama-3-8B-Instruct",
-  "mistral-ai/Ministral-3B": "Ministral 3B",
-  "mistral-ai/Mistral-Large-2411": "Mistral Large 24.11",
-  "mistral-ai/Mistral-Nemo": "Mistral Nemo",
-  "mistral-ai/Mistral-large-2407": "Mistral Large (2407)",
-  "mistral-ai/Mistral-small": "Mistral Small",
-  "cohere/cohere-command-a": "Cohere Command A",
-  "core42/jais-30b-chat": "JAIS 30b Chat",
-  "mistral-ai/mistral-small-2503": "Mistral Small 3.1"
+    "openai/gpt-4.1": "OpenAI GPT-4.1",
+    "openai/gpt-4.1-mini": "OpenAI GPT-4.1-mini",
+    "deepseek/DeepSeek-R1": "DeepSeek-R1",
+    "microsoft/Phi-3.5-mini-instruct": "Phi-3.5-mini instruct",
+    "meta/Meta-Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct",
+    # Add more as needed...
 }
 
 
-async def generate_ai_response(prompt: str, model: str):
-    # Configuration for unofficial GitHub AI endpoint
+async def generate_ai_response(chat_id: str, prompt: str, model: str):
     token = os.getenv("GITHUB_TOKEN")
     if not token:
         raise HTTPException(status_code=500, detail="GitHub token not configured")
-    
-    endpoint = "https://models.github.ai/inference"
-    
-    # Validate the model
+
     if model not in AVAILABLE_MODELS:
-        raise HTTPException(status_code=400, detail=f"Model not available. Choose from: {', '.join(AVAILABLE_MODELS.keys())}")
+        raise HTTPException(status_code=400, detail=f"Invalid model. Choose from: {', '.join(AVAILABLE_MODELS)}")
 
+    endpoint = "https://models.github.ai/inference"
     client = AsyncOpenAI(base_url=endpoint, api_key=token)
 
+    # Retrieve or initialize message history
+    messages = chat_history.get(chat_id, [])
+    if not messages:
+        messages.append({"role": "system", "content": system})
+
+    messages.append({"role": "user", "content": prompt})
+
     try:
         stream = await client.chat.completions.create(
-            messages=[
-                {"role": "system", "content": system},
-                {"role": "user", "content": prompt}
-            ],
+            messages=messages,
             model=model,
             temperature=1.0,
             top_p=1.0,
             stream=True
         )
 
+        # Update history only if generation starts
+        chat_history[chat_id] = messages
+
         async for chunk in stream:
             if chunk.choices and chunk.choices[0].delta.content:
                 yield chunk.choices[0].delta.content
@@ -106,18 +79,27 @@ async def generate_ai_response(prompt: str, model: str):
         yield f"Error: {str(err)}"
         raise HTTPException(status_code=500, detail="AI generation failed")
 
+
 @app.post("/generate")
 async def generate_response(
-    prompt: str = Query(..., description="The prompt for the AI"),
-    model: str = Query("openai/gpt-4.1-mini", description="The model to use for generation")
+    chat_id: str = Query(..., description="Chat session ID"),
+    prompt: str = Query(..., description="User prompt"),
+    model: str = Query("openai/gpt-4.1-mini", description="Model name")
 ):
     if not prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
-    
+
     return StreamingResponse(
-        generate_ai_response(prompt, model),
+        generate_ai_response(chat_id, prompt, model),
         media_type="text/event-stream"
     )
 
+
+@app.post("/reset")
+async def reset_chat(chat_id: str = Query(..., description="Chat session ID to reset")):
+    chat_history.pop(chat_id, None)
+    return {"message": f"Chat history for {chat_id} has been cleared."}
+
+
 def get_app():
-    return app
+    return app