OpenAPI API QA support (#9)

* chore: Added project urls and fixed python vulnerability scan paths

* chore: Added some fixes and documentation to py_function_parser.  (#2)

* chore: Added some fixes and documentation to py_function_parser. Added py parser tests.

* fix: Fixing vuln scan action

* feat: Added LLM tools, Added class inheritance for classes that need LLM handling.

* feat: Added support for enums.

Updated tests.
Added links to inspiration projects

* feat: First 'naive' iteration of openai function calling with openAPI specs

* chore: version bump

* feat: OpenAPI Parsing and calling is now implemented and working.

* feat: Jinja2 template rendering with prompt and OpenAI Functions

* feat: Added custom filter support.

* feat: Refactored OpenAPI calling.

Added an abstraction for OpenAI functions OpenAIFunctionWrapper. It is still in early dev.

* feat: Added json serialization of wrappers.

* feat: Fixing python vulnerability scan workflow.

* chore: Version bump for json serialization feature of OpenAPI wrapper

* feat: OpenAPI spec chatbot (#8)

Refs: #7

* feat: Added OpenAPI QA bot support and example gradio chatbot

Refs: #7

---------

Signed-off-by: Trayan Azarov <[email protected]>

Author: tazarov

README.md

@@ -119,6 +119,61 @@ Name: John
 Age: 20
 """
 ```
+### OpenAPI Spec Chat Bot
+
+This example starts a `gradio` server that allows you to interact with the OpenAPI spec.
+
+```python
+import gradio as gr
+from dotenv import load_dotenv
+
+from func_ai.utils.llm_tools import OpenAIInterface
+from func_ai.utils.openapi_function_parser import OpenAPISpecOpenAIWrapper
+
+_chat_message = []
+
+_spec = None
+
+
+def add_text(history, text):
+    global _chat_message
+    history = history + [(text, None)]
+    _chat_message.append(_spec.api_qa(text, max_tokens=500))
+    return history, ""
+
+
+def add_file(history, file):
+    history = history + [((file.name,), None)]
+    return history
+
+
+def bot(history):
+    global _chat_message
+    # print(temp_callback_handler.get_output())
+    # response = temp_callback_handler.get_output()['output']
+    history[-1][1] = _chat_message[-1]
+    return history
+
+
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=1500)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            txt = gr.Textbox(
+                show_label=False,
+                placeholder="Enter text and press enter",
+            ).style(container=False)
+    txt.submit(add_text, [chatbot, txt], [chatbot, txt]).then(
+        bot, chatbot, chatbot
+    )
+
+if __name__ == "__main__":
+    load_dotenv()
+    _spec = OpenAPISpecOpenAIWrapper.from_url('http://petstore.swagger.io/v2/swagger.json',
+                                              llm_interface=OpenAIInterface(), index=True)
+    demo.launch()
+```
 
 ## Inspiration
 

func_ai/function_indexer.py

@@ -1,35 +1,40 @@
 import inspect
-import json
+import logging
 
 import chromadb
 from chromadb import Settings
 from chromadb.utils import embedding_functions
 from ulid import ULID
 
+from func_ai.utils.llm_tools import OpenAIFunctionWrapper
 from func_ai.utils.py_function_parser import func_to_json
 
-openai_ef = embedding_functions.OpenAIEmbeddingFunction(
-    model_name="text-embedding-ada-002"
-)
+logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.DEBUG)
+logger = logging.getLogger(__name__)
 
 
 class FunctionIndexer(object):
     """
     Index functions
     """
 
-    def __init__(self, db_path: str) -> None:
+    def __init__(self, db_path: str, collection_name: str = "function_index") -> None:
         """
         Initialize function indexer
-        :param db_path:
+        :param db_path: The path where to store the database
+        :param collection_name: The name of the collection
         """
         self._client = chromadb.Client(Settings(
             chroma_db_impl="duckdb+parquet",
             persist_directory=db_path  # Optional, defaults to .chromadb/ in the current directory
         ))
+        self.collection_name = collection_name
         self._fns_map = {}
         self._fns_index_map = {}
         self._open_ai_function_map = []
+        self.openai_ef = embedding_functions.OpenAIEmbeddingFunction(
+            model_name="text-embedding-ada-002"
+        )
 
     def reset_function_index(self) -> None:
         """
@@ -50,8 +55,8 @@ def index_functions(self, functions: list[callable]) -> None:
         """
 
         _ai_fun_map, _fns_map, _fns_index_map = FunctionIndexer.get_functions(functions)
-        collection = self._client.get_or_create_collection(name="function_index", metadata={"hnsw:space": "cosine"},
-                                                           embedding_function=openai_ef)
+        collection = self._client.get_or_create_collection(name=self.collection_name, metadata={"hnsw:space": "cosine"},
+                                                           embedding_function=self.openai_ef)
         self._fns_map.update(_fns_map)
         self._open_ai_function_map.extend(_ai_fun_map)
         self._fns_index_map.update(_fns_index_map)
@@ -62,6 +67,22 @@ def index_functions(self, functions: list[callable]) -> None:
                                   _fns_index_map.items()],
                        ids=[str(ULID()) for _ in _fns_index_map.values()])
 
+    def index_wrapper_functions(self, functions: list[OpenAIFunctionWrapper]):
+        """
+        Index one or more functions
+        Note: Function uniqueness is not checked in this version
+        :param functions:
+        :return:
+        """
+        collection = self._client.get_or_create_collection(name=self.collection_name,
+                                                           metadata={"hnsw:space": "cosine"},
+                                                           embedding_function=self.openai_ef)
+        # print(f"Docs: {collection.get()}")
+        collection.add(documents=[f.description for f in functions],
+                       metadatas=[{"name": f.name, **f.metadata_dict} for f in
+                                  functions],
+                       ids=[str(ULID()) for _ in functions])
+
     def rehydrate_function_map(self, functions: list[callable]) -> None:
         """
         Rehydrate function map
@@ -92,13 +113,14 @@ def find_functions(self, query: str, max_results: int = 2) -> callable:
         :return:
         """
         _response = []
-        collection = self._client.get_or_create_collection(name="function_index", metadata={"hnsw:space": "cosine"},
-                                                           embedding_function=openai_ef)
-        print(collection.get())
+        collection = self._client.get_or_create_collection(name=self.collection_name,
+                                                           metadata={"hnsw:space": "cosine"},
+                                                           embedding_function=self.openai_ef)
+        # print(collection.get())
         res = collection.query(query_texts=[query], n_results=max_results)
         print(f"Got response for sematic search: {res}")
         for r in res['metadatas'][0]:
-            _response.append(self._fns_map[r['name']])
+            _response.append(r['name'])
         return _response
 
     @staticmethod

func_ai/ui_demos/__init__.py

@@ -0,0 +1,49 @@
+import gradio as gr
+from dotenv import load_dotenv
+
+from func_ai.utils.llm_tools import OpenAIInterface
+from func_ai.utils.openapi_function_parser import OpenAPISpecOpenAIWrapper
+
+_chat_message = []
+
+_spec = None
+
+
+def add_text(history, text):
+    global _chat_message
+    history = history + [(text, None)]
+    _chat_message.append(_spec.api_qa(text, max_tokens=500))
+    return history, ""
+
+
+def add_file(history, file):
+    history = history + [((file.name,), None)]
+    return history
+
+
+def bot(history):
+    global _chat_message
+    # print(temp_callback_handler.get_output())
+    # response = temp_callback_handler.get_output()['output']
+    history[-1][1] = _chat_message[-1]
+    return history
+
+
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=1500)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            txt = gr.Textbox(
+                show_label=False,
+                placeholder="Enter text and press enter",
+            ).style(container=False)
+    txt.submit(add_text, [chatbot, txt], [chatbot, txt]).then(
+        bot, chatbot, chatbot
+    )
+
+if __name__ == "__main__":
+    load_dotenv()
+    _spec = OpenAPISpecOpenAIWrapper.from_url('http://petstore.swagger.io/v2/swagger.json',
+                                              llm_interface=OpenAIInterface(), index=True)
+    demo.launch()

func_ai/ui_demos/api_qa.py

@@ -0,0 +1,7 @@
+You are an API expert. Your goal is to assist the user in using an API that he/she is not familiar with.
+The user will provide commands which you will use to find out information about an API
+Then the user will ask you questions about the API and you will answer them.
+
+Rules:
+1. You will only answer questions about the API
+2. You will keep your output only to the essential information

func_ai/utils/api_qa_system_prompt.txt

@@ -0,0 +1,12 @@
+import inspect
+
+
+def arg_in_func(func, arg_name):
+    # Get the signature of the function
+    signature = inspect.signature(func)
+
+    # Get the parameters of the function from the signature
+    parameters = signature.parameters
+
+    # Check if the arg_name is in the parameters
+    return arg_name in parameters

func_ai/utils/common.py

@@ -15,7 +15,8 @@
 
 from pydantic import BaseModel, Field
 
-from func_ai.utils.py_function_parser import type_mapping
+from func_ai.utils.common import arg_in_func
+from func_ai.utils.py_function_parser import type_mapping, func_to_json
 
 logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.DEBUG)
 logger = logging.getLogger(__name__)
@@ -44,6 +45,13 @@ def get_conversation(self) -> list:
         """
         return self.conversation
 
+    def get_last_message(self) -> Any:
+        """
+        Returns the last message
+        :return:
+        """
+        return self.conversation[-1]
+
 
 class OpenAIConversationStore(ConversationStore):
 
@@ -129,15 +137,28 @@ def get_conversation(self) -> list[any]:
         """
         return self.conversation_store.get_conversation()
 
-    def add_conversation_message(self, message: any) -> "LLMInterface":
+    def add_conversation_message(self, message: any, update_llm: bool = False, **kwargs) -> "LLMInterface":
         """
         Adds a message to the conversation
-        :param message:
+        :param message: The message to add
+        :param update_llm: Whether to update the LLM with the new conversation
+        :param kwargs: Parameters to pass to the API
         :return:
         """
         self.conversation_store.add_message(message)
+        if update_llm:
+            self.update_llm_conversation(**kwargs)
         return self
 
+    def update_llm_conversation(self, **kwargs) -> "LLMInterface":
+        """
+        Sends the updated conversation to the LLM
+
+        :param kwargs: Parameters to pass to the API
+        :return:
+        """
+        raise NotImplementedError
+
 
 class OpenAIInterface(LLMInterface):
     """
@@ -158,12 +179,15 @@ def __init__(self, *args, **kwargs):
 
     @retry(stop=stop_after_attempt(3), reraise=True, wait=wait_fixed(1),
            retry_error_callback=lambda x: logger.warning(x))
-    def send(self, prompt: str, **kwargs) -> dict:
+    def update_llm_conversation(self, **kwargs) -> "OpenAIInterface":
+        """
+        Sends the updated conversation to the LLM
+
+        :param kwargs: Parameters to pass to the API
+        :return:
+        """
         _functions = kwargs.get("functions", None)
         _model = kwargs.get("model", self.model)
-        # print(type(self._conversation_store))
-        self.conversation_store.add_message({"role": "user", "content": prompt})
-        logger.debug(f"Prompt: {prompt}")
         try:
             if _functions:
                 response = openai.ChatCompletion.create(
@@ -199,12 +223,20 @@ def send(self, prompt: str, **kwargs) -> dict:
             self.update_cost(_model, response)
             _response_message = response["choices"][0]["message"]
             self.conversation_store.add_message(_response_message)
-            return _response_message
+            return self
         except Exception as e:
             logger.error(f"Error: {e}")
             traceback.print_exc()
             raise e
 
+    def send(self, prompt: str, **kwargs) -> dict[str, any]:
+        # print(type(self._conversation_store))
+        self.conversation_store.add_message({"role": "user", "content": prompt})
+        logger.debug(f"Prompt: {prompt}")
+        response = self.update_llm_conversation(**kwargs)
+        logger.debug(f"Response: {response}")
+        return self.conversation_store.get_last_message()
+
     def load_cost_mapping(self, file_path: str) -> None:
         with open(file_path) as f:
             self.cost_mapping = json.load(f)
@@ -221,7 +253,6 @@ def update_cost(self, model, api_response) -> None:
         self.usage[model]["total_tokens"] += api_response['usage']['total_tokens']
 
 
-
 class OpenAISchema(BaseModel):
     @classmethod
     @property
@@ -320,7 +351,10 @@ def __init__(self, llm_interface: LLMInterface, name: str, description: str, par
         assert "required" in parameters, "Required field not present in parameters"
         self._parameters = parameters
         assert callable(func) or isinstance(func, functools.partial), "Function must be callable"
-        self.func = functools.partial(func, action=self)
+        if arg_in_func(func, "action"):
+            self.func = functools.partial(func, action=self)
+        else:
+            self.func = func
         self._metadata = kwargs
         self._llm_calls = []
 
@@ -369,6 +403,15 @@ def metadata(self) -> dict[str, any]:
         """
         return self._metadata
 
+    @property
+    def metadata_dict(self) -> dict[str, str]:
+        """
+        Returns the metadata of the function as a
+
+        :return:
+        """
+        return {k: str(v) for k, v in self._metadata.items()}
+
     @property
     def schema(self) -> dict[str, any]:
         """
@@ -457,3 +500,16 @@ def from_prompt(self, prompt: str, **kwargs) -> "OpenAIFunctionWrapper":
         """
         self.from_response(self.llm_interface.send(prompt, functions=[self.schema]))
         return self
+
+    @classmethod
+    def from_python_function(cls, func: callable, llm_interface: LLMInterface, **kwargs) -> "OpenAIFunctionWrapper":
+        """
+        Returns an instance of the class from Python function
+
+        :param func: Python function
+        :param llm_interface: LLM interface
+        :param kwargs: arguments to be passed to the function
+        :return:
+        """
+        _func = func_to_json(func)
+        return cls(llm_interface=llm_interface, func=func, **_func, **kwargs)