|
|
import chromadb |
|
|
from llama_index.core import VectorStoreIndex |
|
|
from llama_index.vector_stores.chroma import ChromaVectorStore |
|
|
from llama_index.core import StorageContext, Settings |
|
|
from llama_index.core.schema import TextNode |
|
|
from openai import OpenAI |
|
|
from modal_client import ModalClient |
|
|
|
|
|
|
|
|
from structlog import get_logger |
|
|
|
|
|
logger = get_logger(__name__) |
|
|
|
|
|
from typing import Any, List |
|
|
|
|
|
|
|
|
from llama_index.core.bridge.pydantic import PrivateAttr |
|
|
from llama_index.core.embeddings import BaseEmbedding |
|
|
|
|
|
|
|
|
class CustomEmbeddings(BaseEmbedding): |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
base_url:str, |
|
|
api_key:str, |
|
|
model_name: str , |
|
|
**kwargs: Any, |
|
|
) -> None: |
|
|
super().__init__(**kwargs) |
|
|
self._client = OpenAI( |
|
|
base_url=base_url, |
|
|
api_key=api_key |
|
|
) |
|
|
self.model_name = model_name |
|
|
|
|
|
|
|
|
@classmethod |
|
|
def class_name(cls) -> str: |
|
|
return "custom" |
|
|
|
|
|
async def _aget_query_embedding(self, query: str) -> List[float]: |
|
|
return self._get_query_embedding(query) |
|
|
|
|
|
async def _aget_text_embedding(self, text: str) -> List[float]: |
|
|
return self._get_text_embedding(text) |
|
|
|
|
|
def _get_query_embedding(self, query: str) -> List[float]: |
|
|
embeddings = self._client.embeddings.create( |
|
|
model=self.model_name, |
|
|
input=[query] |
|
|
).data[0].embedding |
|
|
return embeddings |
|
|
|
|
|
def _get_text_embedding(self, text: str) -> List[float]: |
|
|
embeddings = self._client.embeddings.create( |
|
|
model=self.model_name, |
|
|
input=[text] |
|
|
).data[0].embedding |
|
|
return embeddings |
|
|
|
|
|
def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]: |
|
|
embeddings_data = self._client.embeddings.create( |
|
|
model=self.model_name, |
|
|
input=texts |
|
|
) |
|
|
return [embedding.embedding for embedding in embeddings_data.data] |
|
|
|
|
|
class EmbeddingService: |
|
|
def __init__(self, collection_name): |
|
|
config = ModalClient.embedding_config() |
|
|
Settings.embed_model = CustomEmbeddings( |
|
|
api_key=config.get("api_key"), |
|
|
base_url=config.get("base_url"), |
|
|
model_name=config.get("model"), |
|
|
embed_batch_size=32 |
|
|
) |
|
|
Settings.chunk_size = 1024 |
|
|
chroma_client = chromadb.EphemeralClient() |
|
|
chroma_collection = chroma_client.create_collection(collection_name) |
|
|
vector_store = ChromaVectorStore(chroma_collection=chroma_collection) |
|
|
self.storage_context = StorageContext.from_defaults(vector_store=vector_store) |
|
|
|
|
|
def prepare_index(self,nodes): |
|
|
|
|
|
self.index = VectorStoreIndex.from_documents(nodes, storage_context=self.storage_context) |
|
|
|
|
|
|
|
|
def infer(self, query, top_k=10): |
|
|
retriever = self.index.as_retriever(similarity_top_k=top_k) |
|
|
results = retriever.retrieve(query) |
|
|
text = "" |
|
|
for result in results: |
|
|
text += "\n -------------------------- \n" |
|
|
text += f"name = {result.metadata['name']}\n" |
|
|
text += f"filename = {result.metadata['filename']}\n" |
|
|
text += f"type = {result.metadata['type']}\n" |
|
|
text += f"namespace = {result.metadata['namespace']}\n" |
|
|
text += f"content = {result.text}\n" |
|
|
return text |
|
|
|
|
|
|
|
|
|
