108 lines
4.3 KiB
Python
108 lines
4.3 KiB
Python
<|channel|>final code<|message|>import time
|
||
from typing import Callable, Awaitable
|
||
|
||
import uvicorn
|
||
from fastapi import FastAPI, Request, Response, HTTPException
|
||
from fastapi.middleware.base import BaseHTTPMiddleware
|
||
from langchain.llms.openai import OpenAI
|
||
from qdrant_client import QdrantClient
|
||
from rich.console import Console
|
||
|
||
# Консоль для красивого логирования
|
||
console = Console()
|
||
|
||
# Инициализация клиента Qdrant (предполагается, что он уже запущен)
|
||
qdrant = QdrantClient(host="localhost", port=6333)
|
||
|
||
# Инициализация LLM (замените на свой ключ API)
|
||
llm = OpenAI(api_key="YOUR_OPENAI_API_KEY")
|
||
|
||
app = FastAPI(title="Human‑in‑the‑loop Middleware Demo")
|
||
|
||
|
||
class HumanInLoopMiddleware(BaseHTTPMiddleware):
|
||
"""
|
||
Middleware, который позволяет прерывать и возобновлять обработку запросов.
|
||
Если в заголовке `X-HITL-Interrupt` присутствует значение 'true',
|
||
запрос будет поставлен в очередь на паузу. При получении
|
||
запроса с заголовком `X-HITL-Resume` выполнение продолжается.
|
||
"""
|
||
|
||
def __init__(self, app: FastAPI):
|
||
super().__init__(app)
|
||
# Очередь для хранения приостановленных запросов (id -> request data)
|
||
self.paused_requests = {}
|
||
|
||
async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
|
||
interrupt_flag = request.headers.get("X-HITL-Interrupt", "false").lower()
|
||
resume_flag = request.headers.get("X-HITL-Resume", "false").lower()
|
||
|
||
# Если запрос должен быть приостановлен
|
||
if interrupt_flag == "true":
|
||
req_id = str(time.time())
|
||
self.paused_requests[req_id] = {
|
||
"method": request.method,
|
||
"url": request.url.path,
|
||
"body": await request.body(),
|
||
"headers": dict(request.headers),
|
||
}
|
||
console.log(f"[yellow]Request {req_id} paused[/]")
|
||
return Response(content=f"Request paused with id: {req_id}", status_code=202)
|
||
|
||
# Если запрос должен возобновиться
|
||
if resume_flag == "true":
|
||
req_id = request.query_params.get("resume_id")
|
||
if not req_id or req_id not in self.paused_requests:
|
||
raise HTTPException(status_code=400, detail="Invalid or missing resume_id")
|
||
|
||
paused = self.paused_requests.pop(req_id)
|
||
console.log(f"[green]Resuming request {req_id}[/]")
|
||
# Создаём новый Request объект из сохранённых данных
|
||
new_request = Request(
|
||
scope={
|
||
"type": "http",
|
||
"method": paused["method"],
|
||
"path": paused["url"],
|
||
"headers": [(k.encode(), v.encode()) for k, v in paused["headers"].items()],
|
||
"query_string": b"",
|
||
"client": request.client,
|
||
"server": request.scope.get("server"),
|
||
},
|
||
receive=lambda: {"type": "http.request", "body": paused["body"]},
|
||
)
|
||
return await call_next(new_request)
|
||
|
||
# Нормальная обработка
|
||
response = await call_next(request)
|
||
return response
|
||
|
||
|
||
app.add_middleware(HumanInLoopMiddleware)
|
||
|
||
|
||
@app.post("/process")
|
||
async def process_endpoint(data: dict):
|
||
"""
|
||
Пример эндпоинта, который использует LLM и Qdrant.
|
||
"""
|
||
console.log("[blue]Processing request...[/]")
|
||
# Сохраняем запрос в Qdrant
|
||
qdrant.upsert(
|
||
collection_name="requests",
|
||
points=[
|
||
{
|
||
"id": str(time.time()),
|
||
"vector": [0.1, 0.2], # placeholder vector
|
||
"payload": data,
|
||
}
|
||
],
|
||
)
|
||
# Генерируем ответ через LLM
|
||
prompt = f"User asked: {data.get('question', '')}"
|
||
answer = llm(prompt)
|
||
console.log("[blue]LLM response generated[/]")
|
||
return {"answer": answer}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=8000) |