Spaces:
Running
Running
File size: 4,141 Bytes
e9d730a 6082154 e9d730a 6082154 e9d730a 6082154 e9d730a 6082154 e9d730a 6082154 e9d730a 6082154 e9d730a 6082154 e9d730a 6082154 e9d730a 6082154 e9d730a b953016 e9d730a 6082154 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
# src/utils/llm_utils.py
from fastapi import HTTPException
from typing import Tuple
import asyncio
import logging
from src.llms.openai_llm import OpenAILanguageModel
from src.llms.ollama_llm import OllamaLanguageModel
from src.llms.bert_llm import BERTLanguageModel
from src.llms.falcon_llm import FalconLanguageModel
from src.llms.llama_llm import LlamaLanguageModel
from src.embeddings.huggingface_embedding import HuggingFaceEmbedding
from src.vectorstores.chroma_vectorstore import ChromaVectorStore
from src.vectorstores.chroma_manager import ChromaManager
from src.utils.logger import logger
from config.config import settings
# Global vector store instance for reuse
_vector_store = None
_embedding_model = None
_vs_lock = asyncio.Lock()
def get_llm_instance(provider: str):
"""
Get LLM instance based on provider
Args:
provider (str): Name of the LLM provider
Returns:
BaseLLM: Instance of the LLM
Raises:
ValueError: If provider is not supported
"""
llm_map = {
'openai': lambda: OpenAILanguageModel(api_key=settings.OPENAI_API_KEY),
'ollama': lambda: OllamaLanguageModel(base_url=settings.OLLAMA_BASE_URL),
'bert': lambda: BERTLanguageModel(),
'falcon': lambda: FalconLanguageModel(),
'llama': lambda: LlamaLanguageModel(),
}
if provider not in llm_map:
raise ValueError(f"Unsupported LLM provider: {provider}")
return llm_map[provider]()
async def get_vector_store() -> Tuple[ChromaVectorStore, HuggingFaceEmbedding]:
"""
Get vector store and embedding model instances with proper initialization
Returns:
Tuple[ChromaVectorStore, HuggingFaceEmbedding]:
Vector store and embedding model instances
"""
global _vector_store, _embedding_model, _vs_lock
async with _vs_lock:
if _vector_store is not None and _embedding_model is not None:
return _vector_store, _embedding_model
try:
# Load embedding model
_embedding_model = HuggingFaceEmbedding(
model_name=settings.EMBEDDING_MODEL)
logger.info(f"Loaded embedding model: {settings.EMBEDDING_MODEL}")
# Get ChromaDB client through the manager
try:
client = await ChromaManager.get_client(
persist_directory=settings.CHROMA_PATH,
reset_if_needed=True
)
logger.info("Successfully initialized ChromaDB client")
except Exception as e:
logger.error(f"Error getting ChromaDB client: {str(e)}")
# Try to reset ChromaDB completely
await ChromaManager.reset_chroma(settings.CHROMA_PATH)
client = await ChromaManager.get_client(
persist_directory=settings.CHROMA_PATH
)
logger.info("Recreated ChromaDB client after reset")
# Create and initialize vector store
_vector_store = ChromaVectorStore(
embedding_function=_embedding_model.embed_documents,
persist_directory=settings.CHROMA_PATH,
collection_name="documents",
client=client
)
# Initialize the vector store
await _vector_store.initialize()
logger.info("Vector store successfully initialized")
return _vector_store, _embedding_model
except Exception as e:
logger.error(f"Error initializing vector store: {str(e)}")
raise HTTPException(
status_code=500,
detail=f"Failed to initialize vector store: {str(e)}"
)
async def cleanup_vectorstore():
"""
Cleanup and reset vector store resources
"""
global _vector_store, _embedding_model, _vs_lock
async with _vs_lock:
_vector_store = None
_embedding_model = None
# Force garbage collection
import gc
gc.collect()
# Reset ChromaDB completely
await ChromaManager.reset_chroma(settings.CHROMA_PATH)
|