From b872750444c709ef66a5ff64518e3a83da07debb Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Tue, 14 Oct 2025 20:06:32 -0400 Subject: [PATCH] Only use OpenAI for embedding --- chunker.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/chunker.py b/chunker.py index 7fef7a6..7eeb73a 100644 --- a/chunker.py +++ b/chunker.py @@ -11,8 +11,6 @@ from dotenv import load_dotenv from llm import LLMClient -USE_OPENAI = os.getenv("OPENAI_API_KEY") != None - load_dotenv() ollama_client = Client(host=os.getenv("OLLAMA_HOST", "http://localhost:11434")) @@ -85,25 +83,16 @@ class Chunk: class Chunker: - embedding_fx = OpenAIEmbeddingFunction( - api_key=os.getenv("OPENAI_API_KEY"), - model_name="text-embedding-3-small", - ) - def __init__(self, collection) -> None: self.collection = collection self.llm_client = LLMClient() def embedding_fx(self, inputs): - if self.llm_client.PROVIDER == "openai": - openai_embedding_fx = OpenAIEmbeddingFunction( - api_key=os.getenv("OPENAI_API_KEY"), - model_name="text-embedding-3-small", - ) - return openai_embedding_fx(inputs) - else: - response = ollama_client.embed(model="mxbai-embed-large", input=inputs[0]) - return response["embeddings"] + openai_embedding_fx = OpenAIEmbeddingFunction( + api_key=os.getenv("OPENAI_API_KEY"), + model_name="text-embedding-3-small", + ) + return openai_embedding_fx(inputs) def chunk_document( self,