This commit is contained in:
2025-08-11 21:38:25 +02:00
parent f87834a1b3
commit 87ba009bd7
4 changed files with 21 additions and 42 deletions

View File

@@ -1,29 +1,25 @@
# app/agents/vectorization_agent.py # app/agents/vectorization_agent.py
from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
import chromadb
client = chromadb.PersistentClient(path="./chroma_db")
vector_store = client.get_or_create_collection(name="documents")
text_splitter = RecursiveCharacterTextSplitter( text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000, chunk_size=1000,
chunk_overlap=100, chunk_overlap=100,
) )
def agent_vectorize_and_store(doc_id: str, text: str, category: str, language: str): def agent_vectorize_and_store(
""" doc_id: str,
Agent 4: Vectorizes a document and stores it in ChromaDB. text: str,
""" category: str,
language: str,
embedding_model,
vector_store
):
print(f"--- [Background Task] Starting vectorization (ID: {doc_id})...") print(f"--- [Background Task] Starting vectorization (ID: {doc_id})...")
try: try:
return
chunks = text_splitter.split_text(text) chunks = text_splitter.split_text(text)
if not chunks: if not chunks:
print(f"--- [Background Task] document {doc_id} has no text to vectorize.") print(f"--- [Background task] document is empty, skip vectorization. (ID: {doc_id})")
return return
chunk_ids = [f"{doc_id}_{i}" for i in range(len(chunks))] chunk_ids = [f"{doc_id}_{i}" for i in range(len(chunks))]
@@ -38,6 +34,6 @@ def agent_vectorize_and_store(doc_id: str, text: str, category: str, language: s
documents=chunks, documents=chunks,
metadatas=metadatas metadatas=metadatas
) )
print(f"--- [Background Task] Document {doc_id} vectorized and stored successfully.") print(f"--- [Background Task] Document {doc_id} vectorized")
except Exception as e: except Exception as e:
print(f"--- [background Task] Vectorization failed (ID: {doc_id}): {e}") print(f"--- [Background Task] Document vectorization failed (ID: {doc_id}): {e}")

View File

@@ -1,47 +1,28 @@
# app/core/vector_store.py
import os import os
import chromadb import chromadb
from dotenv import load_dotenv from dotenv import load_dotenv
from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
load_dotenv() load_dotenv()
LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower() LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower()
embedding_model = None embedding_model = None
print(f"--- [Core] Initializing Embeddings with provider: {LLM_PROVIDER} ---") print(f"--- [Core] Initializing Embeddings with provider: {LLM_PROVIDER} ---")
if LLM_PROVIDER == "azure": if LLM_PROVIDER == "azure":
required_vars = [
"AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_API_KEY",
"OPENAI_API_VERSION", "AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"
]
if not all(os.getenv(var) for var in required_vars):
raise ValueError("One or more Azure OpenAI environment variables for embeddings are not set.")
embedding_model = AzureOpenAIEmbeddings( embedding_model = AzureOpenAIEmbeddings(
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"), azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
api_key=os.getenv("AZURE_OPENAI_API_KEY"), api_key=os.getenv("AZURE_OPENAI_API_KEY"),
api_version=os.getenv("OPENAI_API_VERSION"), api_version=os.getenv("OPENAI_API_VERSION"),
azure_deployment=os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"), azure_deployment=os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME"),
) )
elif LLM_PROVIDER == "openai": elif LLM_PROVIDER == "openai":
if not os.getenv("OPENAI_API_KEY"):
raise ValueError("OPENAI_API_KEY is not set for the 'openai' provider.")
embedding_model = OpenAIEmbeddings( embedding_model = OpenAIEmbeddings(
api_key=os.getenv("OPENAI_API_KEY"), api_key=os.getenv("OPENAI_API_KEY"),
model=os.getenv("OPENAI_EMBEDDING_MODEL_NAME", "text-embedding-3-small") model=os.getenv("OPENAI_EMBEDDING_MODEL_NAME", "text-embedding-3-small")
) )
else: else:
raise ValueError(f"Unsupported LLM_PROVIDER: {LLM_PROVIDER}. Please use 'azure' or 'openai'.") raise ValueError(f"Unsupported LLM_PROVIDER: {LLM_PROVIDER}.")
client = chromadb.PersistentClient(path="./chroma_db") client = chromadb.PersistentClient(path="./chroma_db")
vector_store = client.get_or_create_collection( vector_store = client.get_or_create_collection(name="documents")
name="documents",
metadata={"hnsw:space": "cosine"}
)

View File

@@ -5,10 +5,10 @@ from typing import Dict, Any, List
from fastapi.concurrency import run_in_threadpool from fastapi.concurrency import run_in_threadpool
from PIL import Image from PIL import Image
from io import BytesIO from io import BytesIO
from .. import agents from .. import agents
from ..core.pdf_processor import convert_pdf_to_images, image_to_base64_str from ..core.pdf_processor import convert_pdf_to_images, image_to_base64_str
from ..core.ocr import extract_text_from_images from ..core.ocr import extract_text_from_images
from ..core.vector_store import embedding_model, vector_store
# Create an APIRouter instance # Create an APIRouter instance
router = APIRouter( router = APIRouter(
@@ -102,10 +102,12 @@ async def upload_and_process_document(
full_text = await run_in_threadpool(extract_text_from_images, images) full_text = await run_in_threadpool(extract_text_from_images, images)
background_tasks.add_task( background_tasks.add_task(
agents.agent_vectorize_and_store, agents.agent_vectorize_and_store,
doc_id, doc_id=doc_id,
full_text, text=full_text,
category, category=category,
language language=language,
embedding_model=embedding_model,
vector_store=vector_store
) )
print("--- [Main] Vectorization job added to background tasks.") print("--- [Main] Vectorization job added to background tasks.")
@@ -118,4 +120,4 @@ async def upload_and_process_document(
async def get_result(doc_id: str): async def get_result(doc_id: str):
if doc_id in db_results: if doc_id in db_results:
return db_results[doc_id] return db_results[doc_id]
raise HTTPException(status_code=404, detail="Document not found.") raise HTTPException(status_code=404, detail="Document not found.")

Binary file not shown.