vectara-agentic 0.1.14__py3-none-any.whl → 0.1.16__py3-none-any.whl

vectara_agentic/__init__.py

@@ -3,7 +3,7 @@ vectara_agentic package.
 """
 
 # Define the package version
-__version__ = "0.1.14"
+__version__ = "0.1.16"
 
 # Import classes and functions from modules
 # from .module1 import Class1, function1
@@ -12,6 +12,7 @@ __version__ = "0.1.14"
 
 # Any initialization code
 def initialize_package():
+    """print a message when the package is initialized."""
     print(f"Initializing vectara-agentic version {__version__}...")
 
 

vectara_agentic/_callback.py

@@ -66,7 +66,6 @@ class AgentCallbackHandler(BaseCallbackHandler):
 
     def _handle_agent_step(self, payload: dict) -> None:
         """Calls self.fn() with the information about agent step."""
-        print(f"Handling agent step: {payload}")
         if EventPayload.MESSAGES in payload:
             msg = str(payload.get(EventPayload.MESSAGES))
             if self.fn:

vectara_agentic/_observability.py

@@ -1,19 +1,21 @@
+"""
+Observability for Vectara Agentic.
+"""
 import os
 import json
+from typing import Optional, Union
 import pandas as pd
-
 from .types import ObserverType
 
 def setup_observer() -> bool:
     '''
     Setup the observer.
     '''
+    import phoenix as px
+    from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
+    from phoenix.otel import register
     observer = ObserverType(os.getenv("VECTARA_AGENTIC_OBSERVER_TYPE", "NO_OBSERVER"))
     if observer == ObserverType.ARIZE_PHOENIX:
-        import phoenix as px
-        from phoenix.otel import register
-        from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
-
         phoenix_endpoint = os.getenv("PHOENIX_ENDPOINT", None)
         if not phoenix_endpoint:
             px.launch_app()
@@ -21,7 +23,7 @@ def setup_observer() -> bool:
         elif 'app.phoenix.arize.com' in phoenix_endpoint:   # hosted on Arizze
             phoenix_api_key = os.getenv("PHOENIX_API_KEY", None)
             if not phoenix_api_key:
-                raise Exception("Arize Phoenix API key not set. Please set PHOENIX_API_KEY environment variable.")
+                raise ValueError("Arize Phoenix API key not set. Please set PHOENIX_API_KEY environment variable.")
             os.environ["PHOENIX_CLIENT_HEADERS"] = f"api_key={phoenix_api_key}"
             os.environ["PHOENIX_COLLECTOR_ENDPOINT"] = "https://app.phoenix.arize.com"
             tracer_provider = register(endpoint=phoenix_endpoint, project_name="vectara-agentic")
@@ -29,12 +31,11 @@ def setup_observer() -> bool:
             tracer_provider = register(endpoint=phoenix_endpoint, project_name="vectara-agentic")
         LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
         return True
-    else:
-        print("No observer set.")
-        return False
+    print("No observer set.")
+    return False
 
 
-def _extract_fcs_value(output):
+def _extract_fcs_value(output: Union[str, dict]) -> Optional[float]:
     '''
     Extract the FCS value from the output.
     '''
@@ -49,7 +50,7 @@ def _extract_fcs_value(output):
     return None
 
 
-def _find_top_level_parent_id(row, all_spans):
+def _find_top_level_parent_id(row: pd.Series, all_spans: pd.DataFrame) -> Optional[str]:
     '''
     Find the top level parent id for the given span.
     '''
@@ -67,14 +68,13 @@ def _find_top_level_parent_id(row, all_spans):
     return current_id
 
 
-def eval_fcs():
+def eval_fcs() -> None:
     '''
     Evaluate the FCS score for the VectaraQueryEngine._query span.
     '''
+    import phoenix as px
     from phoenix.trace.dsl import SpanQuery
     from phoenix.trace import SpanEvaluations
-    import phoenix as px
-
     query = SpanQuery().select(
         "output.value",
         "parent_id",
@@ -83,8 +83,10 @@ def eval_fcs():
     client = px.Client()
     all_spans = client.query_spans(query, project_name="vectara-agentic")
     vectara_spans = all_spans[all_spans['name'] == 'VectaraQueryEngine._query'].copy()
-    vectara_spans['top_level_parent_id'] = vectara_spans.apply(lambda row: _find_top_level_parent_id(row, all_spans), axis=1)
-    vectara_spans['score'] = vectara_spans['output.value'].apply(lambda x: _extract_fcs_value(x))
+    vectara_spans['top_level_parent_id'] = vectara_spans.apply(
+        lambda row: _find_top_level_parent_id(row, all_spans), axis=1
+    )
+    vectara_spans['score'] = vectara_spans['output.value'].apply(_extract_fcs_value)
 
     vectara_spans.reset_index(inplace=True)
     top_level_spans = vectara_spans.copy()

vectara_agentic/_prompts.py

@@ -7,10 +7,10 @@ GENERAL_INSTRUCTIONS = """
 - Use tools as your main source of information, do not respond without using a tool. Do not respond based on pre-trained knowledge.
 - When using a tool with arguments, simplify the query as much as possible if you use the tool with arguments.
   For example, if the original query is "revenue for apple in 2021", you can use the tool with a query "revenue" with arguments year=2021 and company=apple.
-- If you can't answer the question with the information provided by the tools, try to rephrase the question and call a tool again,
+- If you can't answer the question with the information provided by a tool, try to rephrase the question and call the tool again,
   or break the question into sub-questions and call a tool for each sub-question, then combine the answers to provide a complete response.
   For example if asked "what is the population of France and Germany", you can call the tool twice, once for each country.
-- If a query tool provides citations or references in markdown as part of its response, include the citations in your response.
+- If a query tool provides citations or references in markdown as part of its response, include the references in your response.
 - When providing links in your response, where possible put the name of the website or source of information for the displayed text. Don't just say 'source'.
 - If after retrying you can't get the information or answer the question, respond with "I don't know".
 - Your response should never be the input to a tool, only the output.
@@ -20,6 +20,15 @@ GENERAL_INSTRUCTIONS = """
 - Be very careful to respond only when you are confident the response is accurate and not a hallucination.
 - If including latex equations in the markdown response, make sure the equations are on a separate line and enclosed in double dollar signs.
 - Always respond in the language of the question, and in text (no images, videos or code).
+- Always call the "get_bad_topics" tool to determine the topics you are not allowed to discuss or respond to.
+- If you are provided with database tools use them for analytical queries (such as counting, calculating max, min, average, sum, or other statistics).
+  For each database, the database tools include: x_list_tables, x_load_data, x_describe_tables, and x_load_sample_data, where 'x' in the database name.
+  The x_list_tables tool provides a list of available tables in the x database.
+  Always use the x_describe_tables tool to understand the schema of each table, before you access data from that table.
+  Always use the x_load_sample_data tool to understand the column names, and the unique values in each column, so you can use them in your queries.
+  Some times the user may ask for a specific column value, but the actual value in the table may be different, and you will need to use the correct value.
+- Never call x_load_data to retrieve values from each row in the table.
+- Do not mention table names or database names in your response.
 """
 
 #
@@ -64,10 +73,7 @@ IMPORTANT - FOLLOW THESE INSTRUCTIONS CAREFULLY:
 {INSTRUCTIONS}
 {custom_instructions}
 
-## Input
-The user will specify a task or a question in text.
-
-### Output Format
+## Output Format
 
 Please answer in the same language as the question and use the following format:
 
@@ -94,12 +100,12 @@ At that point, you MUST respond in the one of the following two formats (and do
 
 ```
 Thought: I can answer without using any more tools. I'll use the user's language to answer
-Answer: [your answer here (In the same language as the user's question, and maintain any references/citations)]
+Answer: [your answer here (In the same language as the user's question, and maintain any references)]
 ```
 
 ```
 Thought: I cannot answer the question with the provided tools.
-Answer: [your answer here (In the same language as the user's question, and maintain any references/citations)]
+Answer: [your answer here (In the same language as the user's question)]
 ```
 
 ## Current Conversation

vectara_agentic/agent.py

@@ -6,6 +6,9 @@ import os
 from datetime import date
 import time
 import json
+import logging
+import traceback
+
 import dill
 from dotenv import load_dotenv
 
@@ -21,12 +24,6 @@ from llama_index.core.callbacks.base_handler import BaseCallbackHandler
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.core.memory import ChatMemoryBuffer
 
-import logging
-logger = logging.getLogger('opentelemetry.exporter.otlp.proto.http.trace_exporter')
-logger.setLevel(logging.CRITICAL)
-
-load_dotenv(override=True)
-
 from .types import AgentType, AgentStatusType, LLMRole, ToolType
 from .utils import get_llm, get_tokenizer_for_model
 from ._prompts import REACT_PROMPT_TEMPLATE, GENERAL_PROMPT_TEMPLATE
@@ -34,6 +31,10 @@ from ._callback import AgentCallbackHandler
 from ._observability import setup_observer, eval_fcs
 from .tools import VectaraToolFactory, VectaraTool
 
+logger = logging.getLogger("opentelemetry.exporter.otlp.proto.http.trace_exporter")
+logger.setLevel(logging.CRITICAL)
+
+load_dotenv(override=True)
 
 def _get_prompt(prompt_template: str, topic: str, custom_instructions: str):
     """
@@ -57,15 +58,14 @@ def _get_prompt(prompt_template: str, topic: str, custom_instructions: str):
 
 def _retry_if_exception(exception):
     # Define the condition to retry on certain exceptions
-    return isinstance(
-        exception, (TimeoutError)
-    )
+    return isinstance(exception, (TimeoutError))
 
 
 class Agent:
     """
     Agent class for handling different types of agents and their interactions.
     """
+
     def __init__(
         self,
         tools: list[FunctionTool],
@@ -73,6 +73,7 @@ class Agent:
         custom_instructions: str = "",
         verbose: bool = True,
         update_func: Optional[Callable[[AgentStatusType, str], None]] = None,
+        agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
         agent_type: AgentType = None,
     ) -> None:
         """
@@ -84,35 +85,38 @@ class Agent:
             topic (str, optional): The topic for the agent. Defaults to 'general'.
             custom_instructions (str, optional): Custom instructions for the agent. Defaults to ''.
             verbose (bool, optional): Whether the agent should print its steps. Defaults to True.
-            update_func (Callable): A callback function the code calls on any agent updates.
+            agent_progress_callback (Callable): A callback function the code calls on any agent updates.
+                update_func (Callable): old name for agent_progress_callback. Will be deprecated in future.
+            agent_type (AgentType, optional): The type of agent to be used. Defaults to None.
         """
         self.agent_type = agent_type or AgentType(os.getenv("VECTARA_AGENTIC_AGENT_TYPE", "OPENAI"))
         self.tools = tools
         self.llm = get_llm(LLMRole.MAIN)
         self._custom_instructions = custom_instructions
         self._topic = topic
+        self.agent_progress_callback = agent_progress_callback if agent_progress_callback else update_func
 
         main_tok = get_tokenizer_for_model(role=LLMRole.MAIN)
         self.main_token_counter = TokenCountingHandler(tokenizer=main_tok) if main_tok else None
         tool_tok = get_tokenizer_for_model(role=LLMRole.TOOL)
         self.tool_token_counter = TokenCountingHandler(tokenizer=tool_tok) if tool_tok else None
 
-        callbacks: list[BaseCallbackHandler] = [AgentCallbackHandler(update_func)]
+        callbacks: list[BaseCallbackHandler] = [AgentCallbackHandler(self.agent_progress_callback)]
         if self.main_token_counter:
             callbacks.append(self.main_token_counter)
         if self.tool_token_counter:
             callbacks.append(self.tool_token_counter)
-        callback_manager = CallbackManager(callbacks)   # type: ignore
+        callback_manager = CallbackManager(callbacks)  # type: ignore
         self.llm.callback_manager = callback_manager
         self.verbose = verbose
 
-        memory = ChatMemoryBuffer.from_defaults(token_limit=128000)
+        self.memory = ChatMemoryBuffer.from_defaults(token_limit=128000)
         if self.agent_type == AgentType.REACT:
             prompt = _get_prompt(REACT_PROMPT_TEMPLATE, topic, custom_instructions)
             self.agent = ReActAgent.from_tools(
                 tools=tools,
                 llm=self.llm,
-                memory=memory,
+                memory=self.memory,
                 verbose=verbose,
                 react_chat_formatter=ReActChatFormatter(system_header=prompt),
                 max_iterations=30,
@@ -123,7 +127,7 @@ class Agent:
             self.agent = OpenAIAgent.from_tools(
                 tools=tools,
                 llm=self.llm,
-                memory=memory,
+                memory=self.memory,
                 verbose=verbose,
                 callable_manager=callback_manager,
                 max_function_calls=20,
@@ -134,7 +138,7 @@ class Agent:
                 tools=tools,
                 llm=self.llm,
                 verbose=verbose,
-                callable_manager=callback_manager
+                callable_manager=callback_manager,
             ).as_agent()
         else:
             raise ValueError(f"Unknown agent type: {self.agent_type}")
@@ -145,14 +149,26 @@ class Agent:
             print(f"Failed to set up observer ({e}), ignoring")
             self.observability_enabled = False
 
+    def clear_memory(self) -> None:
+        """
+        Clear the agent's memory.
+        """
+        self.agent.memory.reset()
+
     def __eq__(self, other):
+        """
+        Compare two Agent instances for equality.
+        """
         if not isinstance(other, Agent):
             print(f"Comparison failed: other is not an instance of Agent. (self: {type(self)}, other: {type(other)})")
             return False
 
         # Compare agent_type
         if self.agent_type != other.agent_type:
-            print(f"Comparison failed: agent_type differs. (self.agent_type: {self.agent_type}, other.agent_type: {other.agent_type})")
+            print(
+                f"Comparison failed: agent_type differs. (self.agent_type: {self.agent_type}, "
+                f"other.agent_type: {other.agent_type})"
+            )
             return False
 
         # Compare tools
@@ -167,7 +183,10 @@ class Agent:
 
         # Compare custom_instructions
         if self._custom_instructions != other._custom_instructions:
-            print(f"Comparison failed: custom_instructions differ. (self.custom_instructions: {self._custom_instructions}, other.custom_instructions: {other._custom_instructions})")
+            print(
+                "Comparison failed: custom_instructions differ. (self.custom_instructions: "
+                f"{self._custom_instructions}, other.custom_instructions: {other._custom_instructions})"
+            )
             return False
 
         # Compare verbose
@@ -177,7 +196,10 @@ class Agent:
 
         # Compare agent
         if self.agent.memory.chat_store != other.agent.memory.chat_store:
-            print(f"Comparison failed: agent memory differs. (self.agent: {repr(self.agent.memory.chat_store)}, other.agent: {repr(other.agent.memory.chat_store)})")
+            print(
+                f"Comparison failed: agent memory differs. (self.agent: {repr(self.agent.memory.chat_store)}, "
+                f"other.agent: {repr(other.agent.memory.chat_store)})"
+            )
             return False
 
         # If all comparisons pass
@@ -192,6 +214,7 @@ class Agent:
         custom_instructions: str = "",
         verbose: bool = True,
         update_func: Optional[Callable[[AgentStatusType, str], None]] = None,
+        agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
         agent_type: AgentType = None,
     ) -> "Agent":
         """
@@ -203,13 +226,18 @@ class Agent:
             topic (str, optional): The topic for the agent. Defaults to 'general'.
             custom_instructions (str, optional): custom instructions for the agent. Defaults to ''.
             verbose (bool, optional): Whether the agent should print its steps. Defaults to True.
-            update_func (Callable): A callback function the code calls on any agent updates.
-
+            agent_progress_callback (Callable): A callback function the code calls on any agent updates.
+                update_func (Callable): old name for agent_progress_callback. Will be deprecated in future.
+            agent_type (AgentType, optional): The type of agent to be used. Defaults to None.
 
         Returns:
             Agent: An instance of the Agent class.
         """
-        return cls(tools, topic, custom_instructions, verbose, update_func, agent_type)
+        return cls(
+            tools=tools, topic=topic, custom_instructions=custom_instructions,
+            verbose=verbose, agent_progress_callback=agent_progress_callback,
+            update_func=update_func, agent_type=agent_type
+        )
 
     @classmethod
     def from_corpus(
@@ -220,6 +248,7 @@ class Agent:
         vectara_customer_id: str = str(os.environ.get("VECTARA_CUSTOMER_ID", "")),
         vectara_corpus_id: str = str(os.environ.get("VECTARA_CORPUS_ID", "")),
         vectara_api_key: str = str(os.environ.get("VECTARA_API_KEY", "")),
+        agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
         verbose: bool = False,
         vectara_filter_fields: list[dict] = [],
         vectara_lambda_val: float = 0.005,
@@ -238,10 +267,12 @@ class Agent:
             vectara_customer_id (str): The Vectara customer ID.
             vectara_corpus_id (str): The Vectara corpus ID (or comma separated list of IDs).
             vectara_api_key (str): The Vectara API key.
+            agent_progress_callback (Callable): A callback function the code calls on any agent updates.
             data_description (str): The description of the data.
             assistant_specialty (str): The specialty of the assistant.
             verbose (bool, optional): Whether to print verbose output.
-            vectara_filter_fields (List[dict], optional): The filterable attributes (each dict maps field name to Tuple[type, description]).
+            vectara_filter_fields (List[dict], optional): The filterable attributes
+                (each dict maps field name to Tuple[type, description]).
             vectara_lambda_val (float, optional): The lambda value for Vectara hybrid search.
             vectara_reranker (str, optional): The Vectara reranker name (default "mmr")
             vectara_rerank_k (int, optional): The number of results to use with reranking.
@@ -253,18 +284,19 @@ class Agent:
         Returns:
             Agent: An instance of the Agent class.
         """
-        vec_factory = VectaraToolFactory(vectara_api_key=vectara_api_key,
-                                         vectara_customer_id=vectara_customer_id,
-                                         vectara_corpus_id=vectara_corpus_id)
+        vec_factory = VectaraToolFactory(
+            vectara_api_key=vectara_api_key,
+            vectara_customer_id=vectara_customer_id,
+            vectara_corpus_id=vectara_corpus_id,
+        )
         field_definitions = {}
-        field_definitions['query'] = (str, Field(description="The user query"))  # type: ignore
+        field_definitions["query"] = (str, Field(description="The user query"))  # type: ignore
         for field in vectara_filter_fields:
-            field_definitions[field['name']] = (eval(field['type']),
-                                                Field(description=field['description']))  # type: ignore
-        QueryArgs = create_model(   # type: ignore
-            "QueryArgs",
-            **field_definitions
-        )
+            field_definitions[field["name"]] = (
+                eval(field["type"]),
+                Field(description=field["description"]),
+            )  # type: ignore
+        query_args = create_model("QueryArgs", **field_definitions)  # type: ignore
 
         vectara_tool = vec_factory.create_rag_tool(
             tool_name=tool_name or f"vectara_{vectara_corpus_id}",
@@ -272,8 +304,9 @@ class Agent:
             Given a user query,
             returns a response (str) to a user question about {data_description}.
             """,
-            tool_args_schema=QueryArgs,
-            reranker=vectara_reranker, rerank_k=vectara_rerank_k,
+            tool_args_schema=query_args,
+            reranker=vectara_reranker,
+            rerank_k=vectara_rerank_k,
             n_sentences_before=vectara_n_sentences_before,
             n_sentences_after=vectara_n_sentences_after,
             lambda_val=vectara_lambda_val,
@@ -293,7 +326,7 @@ class Agent:
             topic=assistant_specialty,
             custom_instructions=assistant_instructions,
             verbose=verbose,
-            update_func=None
+            agent_progress_callback=agent_progress_callback,
         )
 
     def report(self) -> None:
@@ -308,7 +341,7 @@ class Agent:
         print(f"Topic = {self._topic}")
         print("Tools:")
         for tool in self.tools:
-            print(f"- {tool._metadata.name}")
+            print(f"- {tool.metadata.name}")
         print(f"Agent LLM = {get_llm(LLMRole.MAIN).metadata.model_name}")
         print(f"Tool LLM = {get_llm(LLMRole.TOOL).metadata.model_name}")
 
@@ -349,7 +382,6 @@ class Agent:
                 eval_fcs()
             return agent_response.response
         except Exception as e:
-            import traceback
             return f"Vectara Agentic: encountered an exception ({e}) at ({traceback.format_exc()}), and can't respond."
 
     # Serialization methods
@@ -371,17 +403,21 @@ class Agent:
             # Serialize each tool's metadata, function, and dynamic model schema (QueryArgs)
             tool_dict = {
                 "tool_type": tool.tool_type.value,
-                "name": tool._metadata.name,
-                "description": tool._metadata.description,
-                "fn": dill.dumps(tool.fn).decode('latin-1') if tool.fn else None,  # Serialize fn
-                "async_fn": dill.dumps(tool.async_fn).decode('latin-1') if tool.async_fn else None,  # Serialize async_fn
-                "fn_schema": tool._metadata.fn_schema.model_json_schema() if hasattr(tool._metadata, 'fn_schema') else None,  # Serialize schema if available
+                "name": tool.metadata.name,
+                "description": tool.metadata.description,
+                "fn": dill.dumps(tool.fn).decode("latin-1") if tool.fn else None,  # Serialize fn
+                "async_fn": dill.dumps(tool.async_fn).decode("latin-1")
+                if tool.async_fn
+                else None,  # Serialize async_fn
+                "fn_schema": tool.metadata.fn_schema.model_json_schema()
+                if hasattr(tool.metadata, "fn_schema")
+                else None,  # Serialize schema if available
             }
             tool_info.append(tool_dict)
 
         return {
             "agent_type": self.agent_type.value,
-            "memory": dill.dumps(self.agent.memory).decode('latin-1'),
+            "memory": dill.dumps(self.agent.memory).decode("latin-1"),
             "tools": tool_info,
             "topic": self._topic,
             "custom_instructions": self._custom_instructions,
@@ -394,13 +430,13 @@ class Agent:
         agent_type = AgentType(data["agent_type"])
         tools = []
 
-        JSON_TYPE_TO_PYTHON = {
-            "string": "str",
-            "integer": "int",
-            "boolean": "bool",
-            "array": "list",
-            "object": "dict",
-            "number": "float",
+        json_type_to_python = {
+            "string": str,
+            "integer": int,
+            "boolean": bool,
+            "array": list,
+            "object": dict,
+            "number": float,
         }
 
         for tool_data in data["tools"]:
@@ -408,29 +444,33 @@ class Agent:
             if tool_data.get("fn_schema"):
                 field_definitions = {}
                 for field, values in tool_data["fn_schema"]["properties"].items():
-                    if 'default' in values:
-                        field_definitions[field] = (eval(JSON_TYPE_TO_PYTHON.get(values['type'], values['type'])),
-                                                    Field(description=values['description'], default=values['default']))  # type: ignore
+                    if "default" in values:
+                        field_definitions[field] = (
+                            json_type_to_python.get(values["type"], values["type"]),
+                            Field(
+                                description=values["description"],
+                                default=values["default"],
+                            ),
+                        )  # type: ignore
                     else:
-                        field_definitions[field] = (eval(JSON_TYPE_TO_PYTHON.get(values['type'], values['type'])),
-                                                    Field(description=values['description']))    # type: ignore
-                query_args_model = create_model(   # type: ignore
-                    "QueryArgs",
-                    **field_definitions
-                )
+                        field_definitions[field] = (
+                            json_type_to_python.get(values["type"], values["type"]),
+                            Field(description=values["description"]),
+                        )  # type: ignore
+                query_args_model = create_model("QueryArgs", **field_definitions)  # type: ignore
             else:
                 query_args_model = create_model("QueryArgs")
 
-            fn = dill.loads(tool_data["fn"].encode('latin-1')) if tool_data["fn"] else None
-            async_fn = dill.loads(tool_data["async_fn"].encode('latin-1')) if tool_data["async_fn"] else None
+            fn = dill.loads(tool_data["fn"].encode("latin-1")) if tool_data["fn"] else None
+            async_fn = dill.loads(tool_data["async_fn"].encode("latin-1")) if tool_data["async_fn"] else None
 
             tool = VectaraTool.from_defaults(
-                tool_type=ToolType(tool_data["tool_type"]),
                 name=tool_data["name"],
                 description=tool_data["description"],
                 fn=fn,
                 async_fn=async_fn,
-                fn_schema=query_args_model  # Re-assign the recreated dynamic model
+                fn_schema=query_args_model,  # Re-assign the recreated dynamic model
+                tool_type=ToolType(tool_data["tool_type"]),
             )
             tools.append(tool)
 
@@ -441,7 +481,7 @@ class Agent:
             custom_instructions=data["custom_instructions"],
             verbose=data["verbose"],
         )
-        memory = dill.loads(data["memory"].encode('latin-1')) if data.get("memory") else None
+        memory = dill.loads(data["memory"].encode("latin-1")) if data.get("memory") else None
         if memory:
             agent.agent.memory = memory
         return agent

vectara_agentic/agent_endpoint.py

@@ -0,0 +1,63 @@
+"""
+This module contains functions to start the agent behind an API endpoint.
+"""
+import os
+import logging
+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.security.api_key import APIKeyHeader
+from pydantic import BaseModel
+import uvicorn
+
+from .agent import Agent
+
+API_KEY_NAME = "X-API-Key"
+API_KEY = os.getenv("VECTARA_AGENTIC_API_KEY", "dev-api-key")
+api_key_header = APIKeyHeader(name=API_KEY_NAME)
+
+class ChatRequest(BaseModel):
+    """
+    A request model for the chat endpoint.
+    """
+    message: str
+
+
+def create_app(agent: Agent) -> FastAPI:
+    """
+    Create a FastAPI application with a chat endpoint.
+    """
+    app = FastAPI()
+    logger = logging.getLogger("uvicorn.error")
+    logging.basicConfig(level=logging.INFO)
+
+    @app.get("/chat", summary="Chat with the agent")
+    async def chat(message: str, api_key: str = Depends(api_key_header)):
+        logger.info(f"Received message: {message}")
+        if api_key != API_KEY:
+            logger.warning("Unauthorized access attempt")
+            raise HTTPException(status_code=403, detail="Unauthorized")
+
+        if not message:
+            logger.error("No message provided in the request")
+            raise HTTPException(status_code=400, detail="No message provided")
+
+        try:
+            response = agent.chat(message)
+            logger.info(f"Generated response: {response}")
+            return {"response": response}
+        except Exception as e:
+            logger.error(f"Error during agent processing: {e}")
+            raise HTTPException(status_code=500, detail="Internal server error") from e
+
+    return app
+
+
+def start_app(agent: Agent, host='0.0.0.0', port=8000):
+    """
+    Start the FastAPI server.
+
+    Args:
+        host (str, optional): The host address for the API. Defaults to '127.0.0.1'.
+        port (int, optional): The port for the API. Defaults to 8000.
+    """
+    app = create_app(agent)
+    uvicorn.run(app, host=host, port=port)