Fix changed to openai embeddings

runtimerevolution · Nov 6, 2024 · 534644b · 534644b
1 parent adf0aee
commit 534644b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/labs/embeddings/vectorizers/chunk_vectorizer.py b/labs/embeddings/vectorizers/chunk_vectorizer.py
@@ -8,7 +8,7 @@
 
 from config import configuration_variables as settings
 from labs.embeddings.base import Embedder
-from labs.embeddings.ollama import OllamaEmbedder
+from labs.embeddings.openai import OpenAIEmbedder
 from labs.embeddings.vectorizers.base import Vectorizer
 
 logger = logging.getLogger(__name__)
@@ -75,7 +75,7 @@ def vectorize_to_database(self, include_file_extensions, repo_destination):
 
         logger.debug("Embedding all repo documents.")
 
-        embedder = Embedder(OllamaEmbedder, model="nomic-embed-text:latest")
+        embedder = Embedder(OpenAIEmbedder)
         embeddings = embedder.embed(prompt=texts)
 
         logger.debug("Storing all embeddings.")

diff --git a/labs/llm.py b/labs/llm.py
@@ -3,7 +3,7 @@
 from config import configuration_variables as settings
 from labs.decorators import time_and_log_function
 from labs.embeddings.base import Embedder
-from labs.embeddings.ollama import OllamaEmbedder
+from labs.embeddings.openai import OpenAIEmbedder
 from labs.embeddings.vectorizers.chunk_vectorizer import ChunkVectorizer
 from labs.litellm_service.local import RequestLocalLLM
 from labs.litellm_service.request import RequestLiteLLM
@@ -138,7 +138,7 @@ def call_llm_with_context(repo_destination, issue_summary):
 
     ChunkVectorizer().vectorize_to_database(None, repo_destination)
     # find_similar_embeddings narrows down codebase to files that matter for the issue at hand.
-    context = Embedder(OllamaEmbedder, model="nomic-embed-text:latest").retrieve_embeddings(issue_summary)
+    context = Embedder(OpenAIEmbedder).retrieve_embeddings(issue_summary)
 
     prompt = get_prompt(issue_summary)
     prepared_context = prepare_context(context, prompt)

diff --git a/labs/tasks/llm.py b/labs/tasks/llm.py
@@ -6,7 +6,7 @@
 from config import configuration_variables as settings
 from labs.celery import app
 from labs.embeddings.base import Embedder
-from labs.embeddings.ollama import OllamaEmbedder
+from labs.embeddings.openai import OpenAIEmbedder
 from labs.embeddings.vectorizers.chunk_vectorizer import ChunkVectorizer
 from labs.llm import get_llm_response, get_prompt, prepare_context
 
@@ -27,7 +27,7 @@ def vectorize_repo_to_database_task(prefix="", repo_destination=""):
 
 @app.task
 def find_similar_embeddings_task(prefix="", issue_body=""):
-    rows = Embedder(OllamaEmbedder, model="nomic-embed-text:latest").retrieve_embeddings(
+    rows = Embedder(OpenAIEmbedder).retrieve_embeddings(
         redis_client.get(f"{prefix}_issue_body") if prefix else issue_body
     )
     similar_embeddings = [(row[0], row[1], row[2]) for row in rows]