llm.py

import os
import httpx
from datetime import datetime, timedelta
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

# LangChain 및 벡터 DB 관련 라이브러리
from langchain_community.llms import LlamaCpp
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
# HuggingFace 임베딩 (로컬에서 CPU/GPU로 동작)
from langchain_huggingface import HuggingFaceEmbeddings

app = FastAPI()

# ----------------------------------------------------
# 1. 로컬 환경 및 AI 엔진 초기화
# ----------------------------------------------------

# (1) 로컬 임베딩 모델 (한국어 성능이 좋은 모델 사용)
embedding_model = HuggingFaceEmbeddings(model_name="jhgan/ko-sroberta-multitask")

# (2) 로컬 Chroma DB 설정 (대화 기록 저장용 수집소)
# 코드 실행 경로에 'conversation_db' 폴더가 생성되며 데이터가 누적됩니다.
vector_db = Chroma(persist_directory="./conversation_db", embedding_function=embedding_model)

# (3) llama.cpp 엔진 초기화
llm = LlamaCpp(
    model_path="./models/llama-3-8b-Instruct-Q4_K_M.gguf", # 내 PC의 GGUF 모델 경로
    n_ctx=4096,
    n_gpu_layers=-1, # GPU 가속 최대화
    streaming=True,
    verbose=False
)

# Fish Speech 로컬 API 주소
FISH_SPEECH_API_URL = "http://localhost:8080/v1/tts"

class ChatRequest(BaseModel):
    question: str

# ----------------------------------------------------
# 2. 핵심 비즈니스 로직 함수
# ----------------------------------------------------

def get_past_context(query: str) -> str:
    """Chroma DB에서 관련된 과거 대화 내역을 날짜와 함께 가져오는 함수"""
    # 유사도 기반으로 과거 대화 문서 3개 검색
    docs = vector_db.similarity_search(query, k=3)

    if not docs:
        return "과거 대화 기록이 없습니다."

    context_chunks = []
    for doc in docs:
        date_str = doc.metadata.get("timestamp", "알 수 없는 날짜")
        context_chunks.append(f"[{date_str}] 대화 내용: {doc.page_content}")

    return "\n".join(context_chunks)

def save_conversation(question: str, answer: str):
    """현재 나눈 대화를 오늘 날짜/시간 메타데이터와 함께 DB에 실시간 저장"""
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    full_text = f"사용자 질문: {question} / AI 답변: {answer}"

    # 텍스트와 메타데이터 구조화
    doc = Document(
        page_content=full_text,
        metadata={"timestamp": current_time}
    )
    # 벡터 DB에 추가 (다음 질문부터 바로 검색 가능)
    vector_db.add_documents([doc])

# ----------------------------------------------------
# 3. FastAPI 엔드포인트 (음성 RAG 스트리밍)
# ----------------------------------------------------

@app.post("/chat-voice-rag")
async def chat_voice_rag_endpoint(request: ChatRequest):

    # 1단계: RAG 가동 - 과거 10일 전 등의 대화 기록을 DB에서 서칭
    past_context = get_past_context(request.question)

    # 2단계: 프롬프트에 과거 기억(Context) 주입
    prompt = f"""너는 사용자와의 과거 대화 기록을 완벽히 기억하고 대화하는 친절한 AI 비서야.
아래 제공되는 과거 대화 기록을 참고해서 사용자의 질문에 자연스럽게 답해줘.

[과거 대화 기록]
{past_context}

[사용자의 현재 질문]
{request.question}

AI 답변:"""

    # 3단계: LLM 답변 생성 및 Fish Speech를 위한 스트리밍 제너레이터 구현
    async def audio_stream_generator():
        full_answer = ""
        sentence_buffer = ""

        # llama.cpp로부터 한 글자씩 비동기로 텍스트를 받아옴
        async for chunk in llm.astream(prompt):
            full_answer += chunk
            sentence_buffer += chunk

            # 음성 연속성을 위해 문장 단위([., ?, !, \n])로 끊어서 Fish Speech에 토스
            if any(char in chunk for char in [".", "?", "!", "\n"]):
                clean_text = sentence_buffer.strip()
                if clean_text:
                    # Fish Speech 로컬 서버에 TTS 요청 전송 (비동기 HTTP 통신)
                    async with httpx.AsyncClient() as client:
                        response = await client.post(
                            FISH_SPEECH_API_URL,
                            json={
                                "text": clean_text,
                                "format": "mp3",          # 오디오 포맷
                                "latency": "normal"       # 지연 시간 옵션
                                # 필요시 여기에 목소리 복제용 reference_audio ID 지정 가능
                            },
                            timeout=10.0
                        )
                        if response.status_code == 200:
                            # Fish Speech가 생성한 오디오 바이너리를 순차적으로 클라이언트에 반환
                            yield response.content

                sentence_buffer = "" # 문장 버퍼 비우기

        # 대화가 완전히 끝난 후, 이 대화 자체를 다시 미래의 RAG를 위해 DB에 누적 저장
        save_conversation(request.question, full_answer)

    # 4단계: 최종 오디오 스트림 리턴 (프론트엔드에서 들어오는 대로 즉시 소리 재생)
    return StreamingResponse(audio_stream_generator(), media_type="audio/mpeg")

if __name__ == "__main__":
    import uvicorn
    # FastAPI 서버는 8000번 포트로 구동합니다.
    uvicorn.run(app, host="0.0.0.0", port=8000)