Update app.py

app.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from typing import List, Optional
 
 app = FastAPI()
 
@@ -14,6 +15,15 @@ class CodeRequest(BaseModel):
     language: str = "python"
     max_tokens: int = 128
 
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatRequest(BaseModel):
+    messages: List[ChatMessage]
+    system: Optional[str] = ""
+    max_tokens: int = 1024
+
 @app.get("/")
 def root():
     return {"status": "DevOS AI is running"}
@@ -22,17 +32,25 @@ def root():
 def complete_code(request: CodeRequest):
     prompt = f"Continue the following {request.language} code:\n{request.code}"
     inputs = tokenizer(prompt, return_tensors="pt")
-    
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=request.max_tokens,
-            temperature=0.2,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    
+        outputs = model.generate(**inputs, max_new_tokens=request.max_tokens,
+            temperature=0.2, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"suggestion": generated[len(prompt):].strip()}
+
+@app.post("/chat")
+def chat(request: ChatRequest):
+    # Build conversation prompt
+    prompt = request.system + "\n\n" if request.system else ""
+    for msg in request.messages[-8:]:  # last 8 messages for context
+        role = "User" if msg.role == "user" else "DevOS AI"
+        prompt += f"{role}: {msg.content}\n"
+    prompt += "DevOS AI:"
+
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=request.max_tokens,
+            temperature=0.4, do_sample=True, pad_token_id=tokenizer.eos_token_id)
     generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    suggestion = generated[len(prompt):]
-    
-    return {"suggestion": suggestion.strip()}
+    reply = generated[len(prompt):].strip()
+    return {"reply": reply}