OpenAI function calling support (#1011)

embedchain/llm/openai.py

@@ -1,7 +1,8 @@
-from typing import Optional
+import json
+from typing import Any, Dict, Optional
 
 from langchain.chat_models import ChatOpenAI
-from langchain.schema import HumanMessage, SystemMessage
+from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
 from embedchain.config import BaseLlmConfig
 from embedchain.helpers.json_serializable import register_deserializable
@@ -10,14 +11,15 @@ from embedchain.llm.base import BaseLlm
 
 @register_deserializable
 class OpenAILlm(BaseLlm):
-    def __init__(self, config: Optional[BaseLlmConfig] = None):
+    def __init__(self, config: Optional[BaseLlmConfig] = None, functions: Optional[Dict[str, Any]] = None):
+        self.functions = functions
         super().__init__(config=config)
 
     def get_llm_model_answer(self, prompt) -> str:
-        response = OpenAILlm._get_answer(prompt, self.config)
+        response = self._get_answer(prompt, self.config)
         return response
 
-    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
+    def _get_answer(self, prompt: str, config: BaseLlmConfig) -> str:
         messages = []
         if config.system_prompt:
             messages.append(SystemMessage(content=config.system_prompt))
@@ -31,11 +33,23 @@ class OpenAILlm(BaseLlm):
         if config.top_p:
             kwargs["model_kwargs"]["top_p"] = config.top_p
         if config.stream:
-            from langchain.callbacks.streaming_stdout import \
-                StreamingStdOutCallbackHandler
+            from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 
             callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
             chat = ChatOpenAI(**kwargs, streaming=config.stream, callbacks=callbacks)
         else:
             chat = ChatOpenAI(**kwargs)
+        if self.functions is not None:
+            from langchain.chains.openai_functions import create_openai_fn_runnable
+            from langchain.prompts import ChatPromptTemplate
+
+            structured_prompt = ChatPromptTemplate.from_messages(messages)
+            runnable = create_openai_fn_runnable(functions=self.functions, prompt=structured_prompt, llm=chat)
+            fn_res = runnable.invoke(
+                {
+                    "input": prompt,
+                }
+            )
+            messages.append(AIMessage(content=json.dumps(fn_res)))
+
         return chat(messages).content