File size: 4,373 Bytes
e5b79b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
import os
import json
import time
from datetime import datetime
from agent.inference import InferenceEngine
from model.config import ModelConfig
import torch
app = FastAPI()
# Enable CORS (Critical for cross-port communication)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Global State
agent = None
CONVERSATIONS_FILE = os.path.join("data", "conversations.json")
# Ensure data directory exists
os.makedirs("data", exist_ok=True)
if not os.path.exists(CONVERSATIONS_FILE):
with open(CONVERSATIONS_FILE, "w") as f:
json.dump({"sessions": {}}, f)
class ChatRequest(BaseModel):
message: str
session_id: str = "default"
mode: str = "text"
def get_agent():
global agent
if agent is None or agent.model is None:
model_path = "sail.pt"
if os.path.exists(model_path):
try:
agent = InferenceEngine(model_path=model_path)
print(f"Agent Model Loaded from {model_path}.")
except Exception as e:
print(f"Failed to load agent: {e}")
return agent
def log_conversation(session_id, role, content):
try:
with open(CONVERSATIONS_FILE, "r") as f:
data = json.load(f)
if session_id not in data["sessions"]:
data["sessions"][session_id] = {
"id": session_id,
"timestamp": datetime.now().isoformat(),
"messages": []
}
data["sessions"][session_id]["messages"].append({
"role": role,
"content": content,
"timestamp": datetime.now().isoformat()
})
# Update preview (first user message)
if role == "user" and len(data["sessions"][session_id]["messages"]) <= 2:
data["sessions"][session_id]["preview"] = content[:50] + ("..." if len(content) > 50 else "")
with open(CONVERSATIONS_FILE, "w") as f:
json.dump(data, f, indent=2)
except Exception as e:
print(f"Failed to log conversation: {e}")
@app.get("/api/chat/status")
async def get_status():
return {"status": "connected", "model": "sail.pt" if os.path.exists("sail.pt") else "none"}
@app.post("/api/chat")
async def chat(request: ChatRequest):
agent_instance = get_agent()
if not agent_instance:
return {"response": "Model not loaded! Please ensure sail.pt exists in the root directory."}
# Log user message
log_conversation(request.session_id, "user", request.message)
# Generate response
response = agent_instance.generate(request.message, max_new_tokens=256, temperature=0.7)
# Log system response
log_conversation(request.session_id, "assistant", response)
return {"response": response}
@app.get("/api/history")
async def get_history():
try:
with open(CONVERSATIONS_FILE, "r") as f:
data = json.load(f)
sessions = []
for s_id, s_data in data["sessions"].items():
sessions.append({
"id": s_id,
"timestamp": s_data["timestamp"],
"preview": s_data.get("preview", "Untitled Chat")
})
# Sort by latest
sessions.sort(key=lambda x: x["timestamp"], reverse=True)
return {"sessions": sessions}
except Exception as e:
return {"sessions": [], "error": str(e)}
@app.get("/api/history/{session_id}")
async def get_session_history(session_id: str):
try:
with open(CONVERSATIONS_FILE, "r") as f:
data = json.load(f)
if session_id in data["sessions"]:
return {"messages": data["sessions"][session_id]["messages"]}
return {"messages": []}
except Exception as e:
return {"messages": [], "error": str(e)}
@app.post("/api/grammar/check")
async def check_grammar(request: dict):
# Simplified placeholder to keep API compatible if needed
return {"status": "success", "analysis": "Grammar check currently disabled in simplified mode."}
|