Samuele E. Locatelli
108 lines
3.8 KiB
Python
108 lines
3.8 KiB
Python
# main.py
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import StreamingResponse
|
|
import redis
|
|
import json
|
|
import httpx
|
|
|
|
r = redis.Redis(host='localhost', port=6379, db=1)
|
|
|
|
def save_chat(user_id, message):
|
|
r.rpush(f"chat:{user_id}", json.dumps(message))
|
|
|
|
def get_chat(user_id):
|
|
messages = r.lrange(f"chat:{user_id}", 0, -1)
|
|
return [json.loads(m.decode('utf-8')) for m in messages]
|
|
|
|
app = FastAPI()
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
LM_STUDIO_URL = "http://10.74.83.100:1234/v1/chat/completions"
|
|
MODEL_NAME = "qwen/qwen3-4b-2507" # update as needed
|
|
|
|
@app.post("/chat")
|
|
async def chat(request: Request):
|
|
data = await request.json()
|
|
user_id = data.get("user_id", "default")
|
|
message = data["message"]
|
|
|
|
save_chat(user_id, {"role": "user", "content": message})
|
|
history = get_chat(user_id)
|
|
|
|
async with httpx.AsyncClient(timeout=None) as client:
|
|
resp = await client.post(LM_STUDIO_URL, json={
|
|
"model": MODEL_NAME,
|
|
"messages": history,
|
|
})
|
|
result = resp.json()
|
|
reply = result["choices"][0]["message"]["content"]
|
|
save_chat(user_id, {"role": "assistant", "content": reply})
|
|
return {"response": reply}
|
|
|
|
@app.post("/chat-stream")
|
|
async def chat_stream(request: Request):
|
|
data = await request.json()
|
|
user_id = data.get("user_id", "default")
|
|
message = data["message"]
|
|
|
|
# Save user message and build history
|
|
save_chat(user_id, {"role": "user", "content": message})
|
|
history = get_chat(user_id)
|
|
|
|
async def event_generator():
|
|
assistant_text = ""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=None) as client:
|
|
async with client.stream("POST", LM_STUDIO_URL, json={
|
|
"model": MODEL_NAME,
|
|
"messages": history,
|
|
"stream": True
|
|
}) as r:
|
|
async for raw_line in r.aiter_lines():
|
|
if not raw_line:
|
|
continue
|
|
# Normalize to standard SSE "data: ..." form
|
|
line = raw_line if raw_line.startswith("data:") else f"data: {raw_line}"
|
|
payload = line[len("data: "):].strip()
|
|
|
|
if payload == "[DONE]":
|
|
# Finalize and flush
|
|
yield "data: [DONE]\n\n"
|
|
break
|
|
|
|
# Echo the SSE line to client
|
|
yield f"data: {payload}\n\n"
|
|
|
|
# Accumulate content for saving to Redis
|
|
try:
|
|
obj = json.loads(payload)
|
|
choice = obj.get("choices", [{}])[0]
|
|
# Handle OpenAI-style streaming objects
|
|
delta = choice.get("delta", {})
|
|
piece = delta.get("content")
|
|
if piece is None:
|
|
# Some servers send "text" instead
|
|
piece = choice.get("text")
|
|
if piece:
|
|
assistant_text += piece
|
|
except Exception:
|
|
# Ignore non-JSON control lines
|
|
pass
|
|
finally:
|
|
if assistant_text:
|
|
save_chat(user_id, {"role": "assistant", "content": assistant_text})
|
|
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no" # helps if behind Nginx
|
|
}
|
|
return StreamingResponse(event_generator(), media_type="text/event-stream", headers=headers)
|