【AI学习路线 12】AI项目实战 - 从零构建完整应用

# config.py
from pydantic_settings import BaseSettings

class Settings(BaseSettings):
    # API配置
    API_TITLE: str = "AI文档助手"
    API_VERSION: str = "1.0.0"
    
    # 模型配置
    OPENAI_API_KEY: str
    LLM_MODEL: str = "gpt-4"
    EMBEDDING_MODEL: str = "text-embedding-ada-002"
    
    # 向量数据库
    CHROMA_PERSIST_DIR: str = "./chroma_db"
    
    # 文档配置
    UPLOAD_DIR: str = "./uploads"
    CHUNK_SIZE: int = 500
    CHUNK_OVERLAP: int = 50
    
    class Config:
        env_file = ".env"

settings = Settings()

# services/document_processor.py
import os
from typing import List
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
from langchain.schema import Document

class DocumentProcessor:
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？", "；", " ", ""]
        )
    
    def load_document(self, file_path: str) -> List[Document]:
        """加载文档"""
        ext = os.path.splitext(file_path)[1].lower()
        
        if ext == '.pdf':
            loader = PyPDFLoader(file_path)
        elif ext in ['.docx', '.doc']:
            loader = Docx2txtLoader(file_path)
        else:
            # 默认作为文本处理
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
            return [Document(page_content=text, metadata={"source": file_path})]
        
        return loader.load()
    
    def split_documents(self, documents: List[Document]) -> List[Document]:
        """切分文档"""
        return self.text_splitter.split_documents(documents)
    
    def process_file(self, file_path: str) -> List[Document]:
        """处理单个文件"""
        documents = self.load_document(file_path)
        return self.split_documents(documents)

# services/rag_service.py
from typing import List, Optional
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory

class RAGService:
    def __init__(self, persist_directory: str, embedding_model: str, llm_model: str):
        self.embeddings = OpenAIEmbeddings(model=embedding_model)
        self.llm = ChatOpenAI(model=llm_model, temperature=0.7)
        
        self.vectorstore = Chroma(
            persist_directory=persist_directory,
            embedding_function=self.embeddings
        )
        
        self.memory = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True
        )
        
        self.prompt_template = """
        你是一个专业的文档助手。请根据以下上下文回答问题。
        如果上下文中没有相关信息，请明确说明，不要编造答案。
        
        上下文：
        {context}
        
        问题：{question}
        
        回答：
        """
    
    def add_documents(self, documents: List) -> None:
        """添加文档到向量库"""
        self.vectorstore.add_documents(documents)
    
    def search(self, query: str, k: int = 4) -> List:
        """检索相关文档"""
        return self.vectorstore.similarity_search(query, k=k)
    
    def query(self, question: str) -> dict:
        """问答"""
        # 检索相关文档
        docs = self.search(question)
        context = "\n\n".join([doc.page_content for doc in docs])
        
        # 构建prompt
        prompt = PromptTemplate.from_template(self.prompt_template)
        
        # 生成回答
        chain = prompt | self.llm
        response = chain.invoke({"context": context, "question": question})
        
        return {
            "answer": response.content,
            "sources": [{"content": doc.page_content[:200], 
                        "source": doc.metadata.get("source", "unknown")} 
                       for doc in docs]
        }
    
    def chat(self, question: str) -> dict:
        """带记忆的对话"""
        # 获取对话历史
        chat_history = self.memory.load_memory_variables({}).get("chat_history", [])
        
        # 检索
        docs = self.search(question)
        context = "\n\n".join([doc.page_content for doc in docs])
        
        # 生成回答
        prompt = PromptTemplate.from_template(self.prompt_template)
        chain = prompt | self.llm
        response = chain.invoke({"context": context, "question": question})
        
        # 保存到记忆
        self.memory.save_context({"input": question}, {"output": response.content})
        
        return {
            "answer": response.content,
            "sources": [{"content": doc.page_content[:200]} for doc in docs]
        }

# routers/documents.py
from fastapi import APIRouter, UploadFile, File, HTTPException
from typing import List
import os
import uuid

router = APIRouter(prefix="/documents", tags=["documents"])

@router.post("/upload")
async def upload_document(file: UploadFile = File(...)):
    """上传文档"""
    # 保存文件
    file_id = str(uuid.uuid4())
    file_path = f"./uploads/{file_id}_{file.filename}"
    
    os.makedirs("./uploads", exist_ok=True)
    
    with open(file_path, "wb") as f:
        content = await file.read()
        f.write(content)
    
    # 处理文档
    from ..services.document_processor import DocumentProcessor
    from ..services.rag_service import RAGService
    from ..config import settings
    
    processor = DocumentProcessor(
        chunk_size=settings.CHUNK_SIZE,
        chunk_overlap=settings.CHUNK_OVERLAP
    )
    
    documents = processor.process_file(file_path)
    
    # 添加到向量库
    rag = RAGService(
        persist_directory=settings.CHROMA_PERSIST_DIR,
        embedding_model=settings.EMBEDDING_MODEL,
        llm_model=settings.LLM_MODEL
    )
    rag.add_documents(documents)
    
    return {
        "file_id": file_id,
        "filename": file.filename,
        "chunks": len(documents)
    }

# routers/chat.py
from fastapi import APIRouter
from pydantic import BaseModel

router = APIRouter(prefix="/chat", tags=["chat"])

class ChatRequest(BaseModel):
    question: str
    use_memory: bool = True

class ChatResponse(BaseModel):
    answer: str
    sources: list

@router.post("/query", response_model=ChatResponse)
async def query(request: ChatRequest):
    """问答接口"""
    from ..services.rag_service import RAGService
    from ..config import settings
    
    rag = RAGService(
        persist_directory=settings.CHROMA_PERSIST_DIR,
        embedding_model=settings.EMBEDDING_MODEL,
        llm_model=settings.LLM_MODEL
    )
    
    if request.use_memory:
        result = rag.chat(request.question)
    else:
        result = rag.query(request.question)
    
    return ChatResponse(**result)

# main.py
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from .config import settings
from .routers import documents, chat

app = FastAPI(
    title=settings.API_TITLE,
    version=settings.API_VERSION
)

# CORS配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 注册路由
app.include_router(documents.router)
app.include_router(chat.router)

@app.get("/")
async def root():
    return {"message": "AI文档助手API", "version": settings.API_VERSION}

@app.get("/health")
async def health():
    return {"status": "healthy"}