vectara-agentic 0.1.23__py3-none-any.whl → 0.1.25__py3-none-any.whl

tests/test_agent.py

@@ -2,7 +2,8 @@ import unittest
 from datetime import date
 
 from vectara_agentic.agent import _get_prompt, Agent, AgentType, FunctionTool
-
+from vectara_agentic.agent_config import AgentConfig
+from vectara_agentic.types import ModelProvider, ObserverType
 
 class TestAgentPackage(unittest.TestCase):
     def test_get_prompt(self):
@@ -43,6 +44,48 @@ class TestAgentPackage(unittest.TestCase):
             "50",
         )
 
+    def test_agent_config(self):
+        def mult(x, y):
+            return x * y
+
+        tools = [
+            FunctionTool.from_defaults(
+                fn=mult, name="mult", description="Multiplication functions"
+            )
+        ]
+        topic = "AI topic"
+        instructions = "Always do as your father tells you, if your mother agrees!"
+        config = AgentConfig(
+            agent_type=AgentType.REACT,
+            main_llm_provider=ModelProvider.ANTHROPIC,
+            main_llm_model_name="claude-3-5-sonnet-20241022",
+            tool_llm_provider=ModelProvider.TOGETHER,
+            tool_llm_model_name="meta-llama/Llama-3.3-70B-Instruct-Turbo",
+            observer=ObserverType.ARIZE_PHOENIX
+        )
+
+        agent = Agent(
+            tools=tools,
+            topic=topic,
+            custom_instructions=instructions,
+            agent_config=config
+        )
+        self.assertEqual(agent.tools, tools)
+        self.assertEqual(agent._topic, topic)
+        self.assertEqual(agent._custom_instructions, instructions)
+        self.assertEqual(agent.agent_type, AgentType.REACT)
+        self.assertEqual(agent.agent_config.observer, ObserverType.ARIZE_PHOENIX)
+        self.assertEqual(agent.agent_config.main_llm_provider, ModelProvider.ANTHROPIC)
+        self.assertEqual(agent.agent_config.tool_llm_provider, ModelProvider.TOGETHER)
+
+        # To run this test, you must have OPENAI_API_KEY in your environment
+        self.assertEqual(
+            agent.chat(
+                "What is 5 times 10. Only give the answer, nothing else"
+            ).replace("$", "\\$"),
+            "50",
+        )
+
     def test_from_corpus(self):
         agent = Agent.from_corpus(
             tool_name="RAG Tool",

tests/test_tools.py

@@ -32,7 +32,7 @@ class TestToolsPackage(unittest.TestCase):
 
         self.assertIsInstance(query_tool, VectaraTool)
         self.assertIsInstance(query_tool, FunctionTool)
-        self.assertEqual(query_tool.tool_type, ToolType.QUERY)
+        self.assertEqual(query_tool.metadata.tool_type, ToolType.QUERY)
 
     def test_tool_factory(self):
         def mult(x, y):
@@ -42,7 +42,7 @@ class TestToolsPackage(unittest.TestCase):
         other_tool = tools_factory.create_tool(mult)
         self.assertIsInstance(other_tool, VectaraTool)
         self.assertIsInstance(other_tool, FunctionTool)
-        self.assertEqual(other_tool.tool_type, ToolType.QUERY)
+        self.assertEqual(other_tool.metadata.tool_type, ToolType.QUERY)
 
     def test_llama_index_tools(self):
         tools_factory = ToolsFactory()
@@ -56,7 +56,7 @@ class TestToolsPackage(unittest.TestCase):
 
         self.assertIsInstance(arxiv_tool, VectaraTool)
         self.assertIsInstance(arxiv_tool, FunctionTool)
-        self.assertEqual(arxiv_tool.tool_type, ToolType.QUERY)
+        self.assertEqual(arxiv_tool.metadata.tool_type, ToolType.QUERY)
 
     def test_public_repo(self):
         vectara_customer_id = "1366999410"

vectara_agentic/__init__.py

@@ -2,22 +2,15 @@
 vectara_agentic package.
 """
 
-# Define the package version
-__version__ = "0.1.23"
-
-# Import classes and functions from modules
-# from .module1 import Class1, function1
-# from .module2 import Class2, function2
-
-
-# Any initialization code
-def initialize_package():
-    """print a message when the package is initialized."""
-    print(f"Initializing vectara-agentic version {__version__}...")
-
-
-initialize_package()
-
-
-# Define the __all__ variable
-# __all__ = ['Class1', 'function1', 'Class2', 'function2']
+from .agent import Agent
+from .tools import VectaraToolFactory, VectaraTool
+
+# Define the __all__ variable for wildcard imports
+__all__ = ['Agent', 'VectaraToolFactory', 'VectaraTool']
+
+# Ensure package version is available
+try:
+    import importlib.metadata
+    __version__ = importlib.metadata.version("vectara_agentic")
+except Exception:
+    __version__ = "0.0.0"  # fallback if not installed

vectara_agentic/_observability.py

@@ -6,16 +6,16 @@ import json
 from typing import Optional, Union
 import pandas as pd
 from .types import ObserverType
+from .agent_config import AgentConfig
 
-def setup_observer() -> bool:
+def setup_observer(config: AgentConfig) -> 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:
+    if config.observer == ObserverType.ARIZE_PHOENIX:
         phoenix_endpoint = os.getenv("PHOENIX_ENDPOINT", None)
         if not phoenix_endpoint:
             px.launch_app()

vectara_agentic/_prompts.py

@@ -5,6 +5,7 @@ This file contains the prompt templates for the different types of agents.
 # General (shared) instructions
 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.
+- Always call the 'get_current_date' tool to ensure you know the exact date when a user asks a question.
 - 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 a tool responds with "I do not have enough information", try one of the following:
@@ -31,7 +32,8 @@ GENERAL_INSTRUCTIONS = """
   - Use the x_load_unique_values tool to understand the unique values in each column.
     Sometimes 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.
   - Use the x_load_sample_data tool to understand the column names, and typical values in each column.
-- For tool arguments that support conditional logic (such as year='>2022'), use only one of these operators: [">=", "<=", "!=", ">", "<", "="].
+- For tool arguments that support conditional logic (such as year='>2022'), use one of these operators: [">=", "<=", "!=", ">", "<", "="],
+  or a range operator, with inclusive or exclusive brackets (such as '[2021,2022]' or '[2021,2023)').
 - Do not mention table names or database names in your response.
 """
 
@@ -42,7 +44,7 @@ GENERAL_PROMPT_TEMPLATE = """
 You are a helpful chatbot in conversation with a user, with expertise in {chat_topic}.
 
 ## Date
-Today's date is {today}.
+Your birth date is {today}.
 
 ## INSTRUCTIONS:
 IMPORTANT - FOLLOW THESE INSTRUCTIONS CAREFULLY:
@@ -62,7 +64,7 @@ You are designed to help with a variety of tasks, from answering questions to pr
 You have expertise in {chat_topic}.
 
 ## Date
-Today's date is {today}.
+Your birth date is {today}.
 
 ## Tools
 You have access to a wide variety of tools.

vectara_agentic/_version.py

@@ -0,0 +1,4 @@
+"""
+Define the version of the package.
+"""
+__version__ = "0.1.25"

vectara_agentic/agent.py

@@ -8,6 +8,7 @@ import time
 import json
 import logging
 import traceback
+import asyncio
 
 import dill
 from dotenv import load_dotenv
@@ -25,12 +26,14 @@ from llama_index.core.callbacks.base_handler import BaseCallbackHandler
 from llama_index.agent.openai import OpenAIAgent
 from llama_index.core.memory import ChatMemoryBuffer
 
-from .types import AgentType, AgentStatusType, LLMRole, ToolType
+from .types import AgentType, AgentStatusType, LLMRole, ToolType, AgentResponse, AgentStreamingResponse
 from .utils import get_llm, get_tokenizer_for_model
 from ._prompts import REACT_PROMPT_TEMPLATE, GENERAL_PROMPT_TEMPLATE, GENERAL_INSTRUCTIONS
 from ._callback import AgentCallbackHandler
 from ._observability import setup_observer, eval_fcs
-from .tools import VectaraToolFactory, VectaraTool
+from .tools import VectaraToolFactory, VectaraTool, ToolsFactory
+from .tools_catalog import get_current_date
+from .agent_config import AgentConfig
 
 logger = logging.getLogger("opentelemetry.exporter.otlp.proto.http.trace_exporter")
 logger.setLevel(logging.CRITICAL)
@@ -91,7 +94,8 @@ class Agent:
         verbose: bool = True,
         update_func: Optional[Callable[[AgentStatusType, str], None]] = None,
         agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
-        agent_type: AgentType = None,
+        query_logging_callback: Optional[Callable[[str, str], None]] = None,
+        agent_config: Optional[AgentConfig] = None,
     ) -> None:
         """
         Initialize the agent with the specified type, tools, topic, and system message.
@@ -104,14 +108,18 @@ class Agent:
             verbose (bool, optional): Whether the agent should print its steps. Defaults to True.
             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.
+            query_logging_callback (Callable): A callback function the code calls upon completion of a query
+            agent_config (AgentConfig, optional): The configuration of the agent.
+                Defaults to AgentConfig(), which reads from environment variables.
         """
-        self.agent_type = agent_type or AgentType(os.getenv("VECTARA_AGENTIC_AGENT_TYPE", "OPENAI"))
-        self.tools = tools
-        self.llm = get_llm(LLMRole.MAIN)
+        self.agent_config = agent_config or AgentConfig()
+        self.agent_type = self.agent_config.agent_type
+        self.tools = tools + [ToolsFactory().create_tool(get_current_date)]
+        self.llm = get_llm(LLMRole.MAIN, config=self.agent_config)
         self._custom_instructions = custom_instructions
         self._topic = topic
         self.agent_progress_callback = agent_progress_callback if agent_progress_callback else update_func
+        self.query_logging_callback = query_logging_callback
 
         main_tok = get_tokenizer_for_model(role=LLMRole.MAIN)
         self.main_token_counter = TokenCountingHandler(tokenizer=main_tok) if main_tok else None
@@ -131,7 +139,7 @@ class Agent:
         if self.agent_type == AgentType.REACT:
             prompt = _get_prompt(REACT_PROMPT_TEMPLATE, topic, custom_instructions)
             self.agent = ReActAgent.from_tools(
-                tools=tools,
+                tools=self.tools,
                 llm=self.llm,
                 memory=self.memory,
                 verbose=verbose,
@@ -142,7 +150,7 @@ class Agent:
         elif self.agent_type == AgentType.OPENAI:
             prompt = _get_prompt(GENERAL_PROMPT_TEMPLATE, topic, custom_instructions)
             self.agent = OpenAIAgent.from_tools(
-                tools=tools,
+                tools=self.tools,
                 llm=self.llm,
                 memory=self.memory,
                 verbose=verbose,
@@ -151,23 +159,24 @@ class Agent:
                 system_prompt=prompt,
             )
         elif self.agent_type == AgentType.LLMCOMPILER:
-            self.agent = LLMCompilerAgentWorker.from_tools(
-                tools=tools,
+            agent_worker = LLMCompilerAgentWorker.from_tools(
+                tools=self.tools,
                 llm=self.llm,
                 verbose=verbose,
                 callable_manager=callback_manager,
-            ).as_agent()
-            self.agent.agent_worker.system_prompt = _get_prompt(
-                _get_llm_compiler_prompt(self.agent.agent_worker.system_prompt, topic, custom_instructions),
+            )
+            agent_worker.system_prompt = _get_prompt(
+                _get_llm_compiler_prompt(agent_worker.system_prompt, topic, custom_instructions),
                 topic, custom_instructions
             )
-            self.agent.agent_worker.system_prompt_replan = _get_prompt(
-                _get_llm_compiler_prompt(self.agent.agent_worker.system_prompt_replan, topic, custom_instructions),
+            agent_worker.system_prompt_replan = _get_prompt(
+                _get_llm_compiler_prompt(agent_worker.system_prompt_replan, topic, custom_instructions),
                 topic, custom_instructions
             )
+            self.agent = agent_worker.as_agent()
         elif self.agent_type == AgentType.LATS:
             agent_worker = LATSAgentWorker.from_tools(
-                tools=tools,
+                tools=self.tools,
                 llm=self.llm,
                 num_expansions=3,
                 max_rollouts=-1,
@@ -181,7 +190,7 @@ class Agent:
             raise ValueError(f"Unknown agent type: {self.agent_type}")
 
         try:
-            self.observability_enabled = setup_observer()
+            self.observability_enabled = setup_observer(self.agent_config)
         except Exception as e:
             print(f"Failed to set up observer ({e}), ignoring")
             self.observability_enabled = False
@@ -252,7 +261,8 @@ class Agent:
         verbose: bool = True,
         update_func: Optional[Callable[[AgentStatusType, str], None]] = None,
         agent_progress_callback: Optional[Callable[[AgentStatusType, str], None]] = None,
-        agent_type: AgentType = None,
+        query_logging_callback: Optional[Callable[[str, str], None]] = None,
+        agent_config: AgentConfig = AgentConfig(),
     ) -> "Agent":
         """
         Create an agent from tools, agent type, and language model.
@@ -265,7 +275,8 @@ class Agent:
             verbose (bool, optional): Whether the agent should print its steps. Defaults to True.
             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.
+            query_logging_callback (Callable): A callback function the code calls upon completion of a query
+            agent_config (AgentConfig, optional): The configuration of the agent.
 
         Returns:
             Agent: An instance of the Agent class.
@@ -273,7 +284,8 @@ class Agent:
         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
+            query_logging_callback=query_logging_callback,
+            update_func=update_func, agent_config=agent_config
         )
 
     @classmethod
@@ -286,6 +298,7 @@ class Agent:
         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,
+        query_logging_callback: Optional[Callable[[str, str], None]] = None,
         verbose: bool = False,
         vectara_filter_fields: list[dict] = [],
         vectara_lambda_val: float = 0.005,
@@ -305,6 +318,7 @@ class Agent:
             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.
+            query_logging_callback (Callable): A callback function the code calls upon completion of a query
             data_description (str): The description of the data.
             assistant_specialty (str): The specialty of the assistant.
             verbose (bool, optional): Whether to print verbose output.
@@ -364,6 +378,7 @@ class Agent:
             custom_instructions=assistant_instructions,
             verbose=verbose,
             agent_progress_callback=agent_progress_callback,
+            query_logging_callback=query_logging_callback,
         )
 
     def report(self) -> None:
@@ -378,7 +393,10 @@ class Agent:
         print(f"Topic = {self._topic}")
         print("Tools:")
         for tool in self.tools:
-            print(f"- {tool.metadata.name}")
+            if hasattr(tool, 'metadata'):
+                print(f"- {tool.metadata.name}")
+            else:
+                print("- tool without metadata")
         print(f"Agent LLM = {get_llm(LLMRole.MAIN).metadata.model_name}")
         print(f"Tool LLM = {get_llm(LLMRole.TOOL).metadata.model_name}")
 
@@ -394,12 +412,32 @@ class Agent:
             "tool token count": self.tool_token_counter.total_llm_token_count if self.tool_token_counter else -1,
         }
 
+    async def _aformat_for_lats(self, prompt, agent_response):
+        llm_prompt = f"""
+        Given the question '{prompt}', and agent response '{agent_response.response}',
+        Please provide a well formatted final response to the query.
+        final response:
+        """
+        agent_response.response = str(self.llm.acomplete(llm_prompt))
+
+    def chat(self, prompt: str) -> AgentResponse:           # type: ignore
+        """
+        Interact with the agent using a chat prompt.
+
+        Args:
+            prompt (str): The chat prompt.
+
+        Returns:
+            AgentResponse: The response from the agent.
+        """
+        return asyncio.run(self.achat(prompt))
+
     @retry(
         retry_on_exception=_retry_if_exception,
         stop_max_attempt_number=3,
         wait_fixed=2000,
     )
-    def chat(self, prompt: str) -> str:
+    async def achat(self, prompt: str) -> AgentResponse:    # type: ignore
         """
         Interact with the agent using a chat prompt.
 
@@ -407,32 +445,79 @@ class Agent:
             prompt (str): The chat prompt.
 
         Returns:
-            str: The response from the agent.
+            AgentResponse: The response from the agent.
         """
 
         try:
             st = time.time()
-            agent_response = self.agent.chat(prompt)
+            agent_response = await self.agent.achat(prompt)
             if self.agent_type == AgentType.LATS:
-                prompt = f"""
-                Given the question '{prompt}', and agent response '{agent_response.response}',
-                Please provide a well formatted final response to the query.
-                final response:
-                """
-                final_response = str(self.llm.complete(prompt))
-            else:
-                final_response = agent_response.response
-
+                await self._aformat_for_lats(prompt, agent_response)
             if self.verbose:
                 print(f"Time taken: {time.time() - st}")
             if self.observability_enabled:
                 eval_fcs()
-            return final_response
+            if self.query_logging_callback:
+                self.query_logging_callback(prompt, agent_response.response)
+            return agent_response
         except Exception as e:
-            return f"Vectara Agentic: encountered an exception ({e}) at ({traceback.format_exc()}), and can't respond."
+            return AgentResponse(
+                response = (
+                    f"Vectara Agentic: encountered an exception ({e}) at ({traceback.format_exc()})"
+                    ", and can't respond."
+                )
+            )
 
-    # Serialization methods
+    def stream_chat(self, prompt: str) -> AgentStreamingResponse:    # type: ignore
+        """
+        Interact with the agent using a chat prompt with streaming.
+        Args:
+            prompt (str): The chat prompt.
+        Returns:
+            AgentStreamingResponse: The streaming response from the agent.
+        """
+        return asyncio.run(self.astream_chat(prompt))
+
+    @retry(
+        retry_on_exception=_retry_if_exception,
+        stop_max_attempt_number=3,
+        wait_fixed=2000,
+    )
+    async def astream_chat(self, prompt: str) -> AgentStreamingResponse:    # type: ignore
+        """
+        Interact with the agent using a chat prompt asynchronously with streaming.
+        Args:
+            prompt (str): The chat prompt.
+        Returns:
+            AgentStreamingResponse: The streaming response from the agent.
+        """
+        try:
+            agent_response = await self.agent.astream_chat(prompt)
+            original_async_response_gen = agent_response.async_response_gen
+
+            # Wrap async_response_gen
+            async def _stream_response_wrapper():
+                async for token in original_async_response_gen():
+                    yield token  # Yield async token to keep streaming behavior
+
+                # After streaming completes, execute additional logic
+                if self.agent_type == AgentType.LATS:
+                    await self._aformat_for_lats(prompt, agent_response)
+                if self.query_logging_callback:
+                    self.query_logging_callback(prompt, agent_response.response)
+                if self.observability_enabled:
+                    eval_fcs()
+
+            agent_response.async_response_gen = _stream_response_wrapper  # Override method
+            return agent_response
+        except Exception as e:
+            raise ValueError(
+                f"Vectara Agentic: encountered an exception ({e}) at ({traceback.format_exc()}), and can't respond."
+            ) from e
 
+    #
+    # Serialization methods
+    #
     def dumps(self) -> str:
         """Serialize the Agent instance to a JSON string."""
         return json.dumps(self.to_dict())
@@ -449,7 +534,7 @@ class Agent:
         for tool in self.tools:
             # Serialize each tool's metadata, function, and dynamic model schema (QueryArgs)
             tool_dict = {
-                "tool_type": tool.tool_type.value,
+                "tool_type": tool.metadata.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
@@ -469,12 +554,13 @@ class Agent:
             "topic": self._topic,
             "custom_instructions": self._custom_instructions,
             "verbose": self.verbose,
+            "agent_config": self.agent_config.to_dict(),
         }
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "Agent":
         """Create an Agent instance from a dictionary."""
-        agent_type = AgentType(data["agent_type"])
+        agent_config = AgentConfig.from_dict(data["agent_config"])
         tools = []
 
         json_type_to_python = {
@@ -523,7 +609,7 @@ class Agent:
 
         agent = cls(
             tools=tools,
-            agent_type=agent_type,
+            agent_config=agent_config,
             topic=data["topic"],
             custom_instructions=data["custom_instructions"],
             verbose=data["verbose"],

vectara_agentic/agent_config.py

@@ -0,0 +1,86 @@
+"""
+Define the AgentConfig dataclass for the Vectara Agentic utilities.
+"""
+
+import os
+from dataclasses import dataclass, field
+from .types import ModelProvider, AgentType, ObserverType
+
+@dataclass(eq=True, frozen=True)
+class AgentConfig:
+    """
+    Centralized configuration for the Vectara Agentic utilities.
+
+    Each field can default to either a hard-coded value or an environment
+    variable. For example, if you have environment variables you want to
+    fall back on, you can default to them here.
+    """
+
+    # Agent type
+    agent_type: AgentType = field(
+        default_factory=lambda: AgentType(
+            os.getenv("VECTARA_AGENTIC_AGENT_TYPE", AgentType.OPENAI.value)
+        )
+    )
+
+    # Main LLM provider & model name
+    main_llm_provider: ModelProvider = field(
+        default_factory=lambda: ModelProvider(
+            os.getenv("VECTARA_AGENTIC_MAIN_LLM_PROVIDER", ModelProvider.OPENAI.value)
+        )
+    )
+
+    main_llm_model_name: str = field(
+        default_factory=lambda: os.getenv("VECTARA_AGENTIC_MAIN_MODEL_NAME", "")
+    )
+
+    # Tool LLM provider & model name
+    tool_llm_provider: ModelProvider = field(
+        default_factory=lambda: ModelProvider(
+            os.getenv("VECTARA_AGENTIC_TOOL_LLM_PROVIDER", ModelProvider.OPENAI.value)
+        )
+    )
+    tool_llm_model_name: str = field(
+        default_factory=lambda: os.getenv("VECTARA_AGENTIC_TOOL_MODEL_NAME", "")
+    )
+
+    # Observer
+    observer: ObserverType = field(
+        default_factory=lambda: ObserverType(
+            os.getenv("VECTARA_AGENTIC_OBSERVER_TYPE", "NO_OBSERVER")
+        )
+    )
+
+    # Endpoint API key
+    endpoint_api_key: str = field(
+        default_factory=lambda: os.getenv("VECTARA_AGENTIC_API_KEY", "dev-api-key")
+    )
+
+    def to_dict(self) -> dict:
+        """
+        Convert the AgentConfig to a dictionary.
+        """
+        return {
+            "agent_type": self.agent_type.value,
+            "main_llm_provider": self.main_llm_provider.value,
+            "main_llm_model_name": self.main_llm_model_name,
+            "tool_llm_provider": self.tool_llm_provider.value,
+            "tool_llm_model_name": self.tool_llm_model_name,
+            "observer": self.observer.value,
+            "endpoint_api_key": self.endpoint_api_key
+        }
+
+    @classmethod
+    def from_dict(cls, config_dict: dict) -> "AgentConfig":
+        """
+        Create an AgentConfig from a dictionary.
+        """
+        return cls(
+            agent_type=AgentType(config_dict["agent_type"]),
+            main_llm_provider=ModelProvider(config_dict["main_llm_provider"]),
+            main_llm_model_name=config_dict["main_llm_model_name"],
+            tool_llm_provider=ModelProvider(config_dict["tool_llm_provider"]),
+            tool_llm_model_name=config_dict["tool_llm_model_name"],
+            observer=ObserverType(config_dict["observer"]),
+            endpoint_api_key=config_dict["endpoint_api_key"]
+        )

vectara_agentic/agent_endpoint.py

@@ -1,7 +1,6 @@
 """
 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
@@ -9,10 +8,9 @@ from pydantic import BaseModel
 import uvicorn
 
 from .agent import Agent
+from .agent_config import AgentConfig
 
-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)
+api_key_header = APIKeyHeader(name="X-API-Key")
 
 class ChatRequest(BaseModel):
     """
@@ -21,18 +19,19 @@ class ChatRequest(BaseModel):
     message: str
 
 
-def create_app(agent: Agent) -> FastAPI:
+def create_app(agent: Agent, config: AgentConfig) -> FastAPI:
     """
     Create a FastAPI application with a chat endpoint.
     """
     app = FastAPI()
     logger = logging.getLogger("uvicorn.error")
     logging.basicConfig(level=logging.INFO)
+    endpoint_api_key = config.endpoint_api_key
 
     @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:
+        if api_key != endpoint_api_key:
             logger.warning("Unauthorized access attempt")
             raise HTTPException(status_code=403, detail="Unauthorized")
 
@@ -59,5 +58,5 @@ def start_app(agent: Agent, host='0.0.0.0', port=8000):
         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)
+    app = create_app(agent, config=AgentConfig())
     uvicorn.run(app, host=host, port=port)