Memory.py

import os
import re
import json
import asyncio
import uuid
import time
import logging
from contextlib import asynccontextmanager
from pydantic import BaseModel, Field
from typing import List, Optional, Dict
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts import PromptTemplate
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
import uvicorn
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.runnables import RunnablePassthrough
from langchain_core.runnables import ConfigurableFieldSpec
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import SQLChatMessageHistory
import yaml
import argparse

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def load_config(config_path):
    """
    加载并解析 YAML 配置文件
    :param config_path: YAML 配置文件路径
    :return: 配置字典
    """
    try:
        with open(config_path, 'r', encoding='utf-8') as file:
            config = yaml.safe_load(file)
        return config
    except FileNotFoundError:
        logger.error(f"错误：配置文件 '{config_path}' 未找到。")
        return None
    except yaml.YAMLError as exc:
        logger.error(f"错误：解析 YAML 文件时出错: {exc}")
        return None

parser = argparse.ArgumentParser(description="读取 YAML 配置文件")
parser.add_argument(
    "--config", 
    type=str, 
    required=True, 
    help="YAML 配置文件的路径"
)
args = parser.parse_args()
config = load_config(args.config)

global API_TYPE,PORT,ONEAPI_CHAT_MODEL
global ONEAPI_API_BASE, ONEAPI_EMBEDDING_API_KEY, ONEAPI_EMBEDDING_MODEL
global CHROMADB_DIRECTORY, CHROMADB_COLLECTION_NAME

PORT = config.get("PORT")
API_TYPE = config.get("API_TYPE")
ONEAPI_CHAT_MODEL = config.get("ONEAPI_CHAT_MODEL")
ONEAPI_API_BASE = config.get("ONEAPI_API_BASE")
ONEAPI_EMBEDDING_MODEL = config.get("ONEAPI_EMBEDDING_MODEL")
ONEAPI_CHAT_API_KEY = config.get("ONEAPI_CHAT_API_KEY")
ONEAPI_EMBEDDING_API_KEY = config.get("ONEAPI_EMBEDDING_API_KEY")
PROMPT_TEMPLATE_TXT = config.get("PROMPT_TEMPLATE_TXT")
CHROMADB_DIRECTORY = config.get("CHROMADB_DIRECTORY")
CHROMADB_COLLECTION_NAME = config.get("CHROMADB_COLLECTION_NAME")
PORT = config.get("PORT")


model = None 
embeddings = None 
vectorstore = None 
prompt = None  
chain = None 
with_message_history = None

class Message(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    messages: List[Message]
    stream: Optional[bool] = False
    userId: Optional[str] = None
    conversationId: Optional[str] = None

class ChatCompletionResponseChoice(BaseModel):
    index: int
    message: Message
    finish_reason: Optional[str] = None

class ChatCompletionResponse(BaseModel):
    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    choices: List[ChatCompletionResponseChoice]
    system_fingerprint: Optional[str] = None


def get_session_history(user_id: str, conversation_id: str):
    return SQLChatMessageHistory(f"{user_id}--{conversation_id}", "sqlite:///memory.db")

def getPrompt(prompt):
    logger.info(f"最后给到LLM的prompt的内容: {prompt}")
    return prompt


def format_response(response):
    paragraphs = re.split(r'\n{2,}', response)
    formatted_paragraphs = []
    for para in paragraphs:
        if '```' in para:
            parts = para.split('```')
            for i, part in enumerate(parts):
                if i % 2 == 1:
                    parts[i] = f"\n```\n{part.strip()}\n```\n"
            para = ''.join(parts)
        else:
            para = para.replace('. ', '.\n')
        formatted_paragraphs.append(para.strip())
    return '\n\n'.join(formatted_paragraphs)


@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, embeddings, vectorstore, prompt, chain,with_message_history
    try:
        logger.info("正在初始化模型、实例化Chroma对象、提取prompt模版、定义chain...")
        
        model = ChatOpenAI(
            base_url=ONEAPI_API_BASE,
            api_key=ONEAPI_CHAT_API_KEY,
            model=ONEAPI_CHAT_MODEL,
            # temperature=0,
            # timeout=None,
            # max_retries=2,
        )
        embeddings = OpenAIEmbeddings(
            base_url=ONEAPI_API_BASE,
            api_key=ONEAPI_EMBEDDING_API_KEY,
            model=ONEAPI_EMBEDDING_MODEL,
            deployment=ONEAPI_EMBEDDING_MODEL
        )    
        
        vectorstore = Chroma(persist_directory=CHROMADB_DIRECTORY,
                             collection_name=CHROMADB_COLLECTION_NAME,
                             embedding_function=embeddings,
                             )
        prompt_template = PromptTemplate.from_file(PROMPT_TEMPLATE_TXT)

        prompt = ChatPromptTemplate.from_messages(
            [
                ("system","你是一个针对健康档案进行问答的机器人。你的任务是根据下述给定的已知信息回答用户问题。"),
                MessagesPlaceholder(variable_name="history"),
                ("human", prompt_template.template)
            ]
        )

        chain = prompt | getPrompt | model
        logger.info("初始化完成")

        with_message_history = RunnableWithMessageHistory(
            chain,
            get_session_history,
            input_messages_key="query",
            history_messages_key="history",
            history_factory_config=[
                ConfigurableFieldSpec(
                    id="user_id",
                    annotation=str,
                    name="User ID",
                    description="Unique identifier for the user.",
                    default="",
                    is_shared=True,
                ),
                ConfigurableFieldSpec(
                    id="conversation_id",
                    annotation=str,
                    name="Conversation ID",
                    description="Unique identifier for the conversation.",
                    default="",
                    is_shared=True,
                ),
            ],
        )
    except Exception as e:
        logger.error(f"初始化过程中出错: {str(e)}")
        raise

    yield
    logger.info("正在关闭...")

app = FastAPI(lifespan=lifespan)

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):

    if not model or not embeddings or not vectorstore or not prompt or not chain:
        logger.error("服务未初始化")
        raise HTTPException(status_code=500, detail="服务未初始化")

    try:
        logger.info(f"收到聊天完成请求: {request}")
        query_prompt = request.messages[-1].content
        logger.info(f"用户问题是: {query_prompt}")

        retriever = vectorstore.similarity_search(
            query=query_prompt,
            k=3,
        )
        result = with_message_history.invoke(
            {"query": query_prompt,"context": retriever},
            config={"configurable": {"user_id": request.userId, "conversation_id": request.conversationId}}
        )

        formatted_response = str(format_response(result.content))
        logger.info(f"格式化的搜索结果: {formatted_response}")

        
        if request.stream:
            async def generate_stream():
                chunk_id = f"chatcmpl-{uuid.uuid4().hex}"
                lines = formatted_response.split('\n')
                for i, line in enumerate(lines):
                    chunk = {
                        "id": chunk_id,
                        "object": "chat.completion.chunk",
                        "created": int(time.time()),
                        "choices": [
                            {
                                "index": 0,
                                "delta": {"content": line + '\n'}, # if i > 0 else {"role": "assistant", "content": ""},
                                "finish_reason": None
                            }
                        ]
                    }
                    yield f"{json.dumps(chunk)}\n"
                    await asyncio.sleep(0.5)
                final_chunk = {
                    "id": chunk_id,
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "choices": [
                        {
                            "index": 0,
                            "delta": {},
                            "finish_reason": "stop"
                        }
                    ]
                }
                yield f"{json.dumps(final_chunk)}\n"

            return StreamingResponse(generate_stream(), media_type="text/event-stream")
        else:
            response = ChatCompletionResponse(
                choices=[
                    ChatCompletionResponseChoice(
                        index=0,
                        message=Message(role="assistant", content=formatted_response),
                        finish_reason="stop"
                    )
                ]
            )
            logger.info(f"发送响应内容: \n{response}")
            return JSONResponse(content=response.model_dump())
    except Exception as e:
        logger.error(f"处理聊天完成时出错:\n\n {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))



if __name__ == "__main__":
    logger.info(f"在端口 {PORT} 上启动服务器")
    uvicorn.run(app, host="0.0.0.0", port=PORT)