Try get rid of OpenAI default embeddings in various LlamaIndex objects

ShoggothAI · Oct 25, 2024 · a9e6340 · a9e6340
1 parent 1e65cdf
commit a9e6340
Show file tree

Hide file tree

Showing 10 changed files with 12 additions and 97 deletions.
diff --git a/examples/data/groupchat/fetch_arxiv_gpt4.py b/examples/data/groupchat/fetch_arxiv_gpt4.py
diff --git a/examples/data/groupchat/fetch_latest_gpt4_paper.py b/examples/data/groupchat/fetch_latest_gpt4_paper.py
diff --git a/examples/data/research_agent_storage/default__vector_store.json b/examples/data/research_agent_storage/default__vector_store.json
diff --git a/examples/data/research_agent_storage/docstore.json b/examples/data/research_agent_storage/docstore.json
diff --git a/examples/data/research_agent_storage/graph_store.json b/examples/data/research_agent_storage/graph_store.json
diff --git a/examples/data/research_agent_storage/image__vector_store.json b/examples/data/research_agent_storage/image__vector_store.json
diff --git a/examples/data/research_agent_storage/index_store.json b/examples/data/research_agent_storage/index_store.json
diff --git a/motleycrew/applications/research_agent/question_prioritizer.py b/motleycrew/applications/research_agent/question_prioritizer.py
@@ -1,7 +1,10 @@
+from typing import Optional
+
 from langchain.prompts import PromptTemplate
 from langchain_core.prompts.base import BasePromptTemplate
 from langchain_core.runnables import RunnableLambda, RunnablePassthrough, chain
 from langchain_core.tools import StructuredTool
+from langchain_core.language_models import BaseLanguageModel
 from pydantic import BaseModel, Field
 
 from motleycrew.applications.research_agent.question import Question
@@ -16,8 +19,9 @@ class QuestionPrioritizerTool(MotleyTool):
     def __init__(
         self,
         prompt: str | BasePromptTemplate = None,
+        llm: Optional[BaseLanguageModel] = None,
     ):
-        langchain_tool = create_question_prioritizer_langchain_tool(prompt=prompt)
+        langchain_tool = create_question_prioritizer_langchain_tool(prompt=prompt, llm=llm)
 
         super().__init__(langchain_tool)
 
@@ -47,6 +51,7 @@ class QuestionPrioritizerInput(BaseModel, arbitrary_types_allowed=True):
 
 def create_question_prioritizer_langchain_tool(
     prompt: str | BasePromptTemplate = None,
+    llm: Optional[BaseLanguageModel] = None,
 ) -> StructuredTool:
     if prompt is None:
         prompt = _default_prompt
@@ -56,6 +61,7 @@ def create_question_prioritizer_langchain_tool(
         name="Question prioritizer",
         description="Takes the original question and a list of derived questions, "
         "and selects from the latter the one most pertinent to the former",
+        llm=llm,
     )
 
     @chain

diff --git a/motleycrew/applications/research_agent/question_task.py b/motleycrew/applications/research_agent/question_task.py
@@ -41,7 +41,7 @@ def __init__(
         self.n_iter = 0
         self.question = Question(question=question)
         self.graph_store.insert_node(self.question)
-        self.question_prioritization_tool = QuestionPrioritizerTool()
+        self.question_prioritization_tool = QuestionPrioritizerTool(llm=llm)
         self.question_generation_tool = QuestionGeneratorTool(
             query_tool=query_tool, graph=self.graph_store, llm=llm
         )

diff --git a/motleycrew/tools/simple_retriever_tool.py b/motleycrew/tools/simple_retriever_tool.py
@@ -65,14 +65,16 @@ def make_retriever_langchain_tool(
         # load the documents and create the index
         documents = SimpleDirectoryReader(data_dir).load_data()
         index = VectorStoreIndex.from_documents(
-            documents, transformations=[SentenceSplitter(chunk_size=512), embeddings]
+            documents,
+            transformations=[SentenceSplitter(chunk_size=512), embeddings],
+            embed_model=embeddings,
         )
         # store it for later
         index.storage_context.persist(persist_dir=persist_dir)
     else:
         # load the existing index
         storage_context = StorageContext.from_defaults(persist_dir=persist_dir)
-        index = load_index_from_storage(storage_context)
+        index = load_index_from_storage(storage_context, embed_model=embeddings)
 
     retriever = index.as_retriever(
         similarity_top_k=10,