src/agent_server/main.py - master - man4j/agent-server

src/agent_server/main.py
91 строка2 KB
Vladimir
fixes
19 апр 2026, 23:14
19 апр 2026, 23:14d45c36c
from chainlit.types import ThreadDict
from chainlit.utils import mount_chainlit
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
import uvicorn

import chainlit as cl
from agent_server.chat.runtime import ChatRuntime
from agent_server.chat.chainlit_context import get_auth_users_map_cached
from agent_server.migrations_runner import run_sql_migrations
from agent_server.pg_storage import PostgresStorageClient
from agent_server.auth.passwords import verify_password
from agent_server.config.env import BASE_URL, DATABASE_URL


app = FastAPI()


_storage_client = PostgresStorageClient(
    dsn=DATABASE_URL,  # тут обычный postgres dsn, не sqlalchemy async url
    public_base_url=BASE_URL,
    route_prefix="/storage",
)
_chat_runtime = ChatRuntime(_storage_client)


@app.on_event("shutdown")
async def shutdown_storage_client():
    await _storage_client.close()


@app.get("/storage/{object_key:path}")
async def get_storage_object(object_key: str):
    item = await _storage_client.fetch_file(object_key)
    if item is None:
        raise HTTPException(status_code=404, detail="Object not found")
    content, mime = item
    return Response(content=content, media_type=mime)

@cl.data_layer
def get_data_layer():
    return _chat_runtime.get_data_layer()


@cl.password_auth_callback
def auth_callback(username: str, password: str):
    auth_users_map = get_auth_users_map_cached()
    user = auth_users_map.get(username)

    if user and verify_password(password, user.password_hash):
        return cl.User(
            identifier=user.username,
            metadata={"role": user.role, "provider": user.provider},
        )

    return None


@cl.set_chat_profiles
async def chat_profile(current_user: cl.User):
    return await _chat_runtime.get_chat_profiles(current_user)


@cl.on_chat_start
async def on_chat_start():
    await _chat_runtime.on_chat_start()


@cl.on_chat_resume
async def on_chat_resume(thread: ThreadDict):
    await _chat_runtime.on_chat_resume(thread)


@cl.on_chat_end
async def on_chat_end():
    await _chat_runtime.on_chat_end()


@cl.on_message
async def on_message(input_msg: cl.Message):
    await _chat_runtime.on_message(input_msg)


def main():
    run_sql_migrations()
    mount_chainlit(app=app, target=__file__, path="/")
    uvicorn.run(app, host="0.0.0.0", port=8000)


if __name__ == "__main__":
    main()