add defualt LLM to chat engine

pinecone-io · Sep 14, 2023 · 690f813 · 690f813
1 parent 026793e
commit 690f813
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/context_engine/chat_engine/chat_engine.py b/context_engine/chat_engine/chat_engine.py
@@ -8,7 +8,7 @@
                                                         FunctionCallingQueryGenerator, )
 from context_engine.context_engine import ContextEngine
 from context_engine.tokenizer import Tokenizer
-from context_engine.llm import BaseLLM
+from context_engine.llm import BaseLLM, OpenAILLM
 from context_engine.llm.models import ModelParams, SystemMessage
 from context_engine.models.api_models import (StreamingChatChunk, ChatResponse,
                                               StreamingChatResponse, )
@@ -54,11 +54,12 @@ async def aget_context(self, messages: Messages) -> Context:
 class ChatEngine(BaseChatEngine):
 
     DEFAULT_QUERY_GENERATOR = FunctionCallingQueryGenerator
+    DEFAULT_LLM = OpenAILLM
 
     def __init__(self,
                  *,
-                 llm: BaseLLM,
                  context_engine: ContextEngine,
+                 llm: Optional[BaseLLM] = None,
                  max_prompt_tokens: int = 4096,
                  max_generated_tokens: Optional[int] = None,
                  max_context_tokens: Optional[int] = None,
@@ -67,7 +68,7 @@ def __init__(self,
                  history_pruning: str = "recent",
                  min_history_messages: int = 1
                  ):
-        self.llm = llm
+        self.llm = llm if llm is not None else self.DEFAULT_LLM()
         self.context_engine = context_engine
         self.max_prompt_tokens = max_prompt_tokens
         self.max_generated_tokens = max_generated_tokens

diff --git a/service/app.py b/service/app.py
@@ -174,7 +174,7 @@ def _init_engines():
     context_engine = ContextEngine(knowledge_base=kb)
     llm = OpenAILLM()
 
-    chat_engine = ChatEngine(llm=llm, context_engine=context_engine)
+    chat_engine = ChatEngine(context_engine=context_engine, llm=llm)
 
 
 def start():

diff --git a/tests/unit/chat_engine/test_chat_engine.py b/tests/unit/chat_engine/test_chat_engine.py
@@ -31,8 +31,8 @@ def _init_chat_engine(self,
                           max_generated_tokens: int = 200,
                           **kwargs):
         return ChatEngine(
-            llm=self.mock_llm,
             context_engine=self.mock_context_engine,
+            llm=self.mock_llm,
             query_builder=self.mock_query_builder,
             system_prompt=system_prompt,
             max_prompt_tokens=max_prompt_tokens,