RDBC-1059 Add AiAgentParameter Policy and ValueType

Author: poissoncorp

ravendb/__init__.py

@@ -91,6 +91,8 @@
     AiAgentConfiguration,
     AiAgentConfigurationResult,
     AiAgentParameter,
+    AiAgentParameterPolicy,
+    AiAgentParameterValueType,
     AiAgentToolAction,
     AiAgentToolQuery,
     AiAgentPersistenceConfiguration,

ravendb/documents/operations/ai/agents/__init__.py

@@ -1,6 +1,8 @@
 from .ai_agent_configuration import (
     AiAgentConfiguration,
     AiAgentParameter,
+    AiAgentParameterPolicy,
+    AiAgentParameterValueType,
     AiAgentToolAction,
     AiAgentToolQuery,
     AiAgentToolQueryOptions,
@@ -38,6 +40,8 @@
     "AiAgentConfiguration",
     "AiAgentConfigurationResult",
     "AiAgentParameter",
+    "AiAgentParameterPolicy",
+    "AiAgentParameterValueType",
     "AiAgentToolAction",
     "AiAgentToolQuery",
     "AiAgentToolQueryOptions",

ravendb/documents/operations/ai/agents/ai_agent_configuration.py

@@ -1,58 +1,79 @@
 from __future__ import annotations
+import enum
 from typing import List, Optional, Dict, Any, Union
 
 
-class AiAgentParameter:
-    """
-    Represents a parameter for an AI agent configuration.
-    Parameters can be used to pass values to the agent's system prompt.
-    """
+class AiAgentParameterPolicy(enum.IntFlag):
+    # FORBID_MODEL_GENERATION blocks a parent agent from generating values for
+    # this parameter when invoking a sub-agent that declares it.
+    DEFAULT = 0
+    FORBID_MODEL_GENERATION = 1
+
+
+class AiAgentParameterValueType(enum.Enum):
+    DEFAULT = "Default"
+    STRING = "String"
+    NUMBER = "Number"
+    BOOLEAN = "Boolean"
+    ARRAY_OF_STRING = "ArrayOfString"
+    ARRAY_OF_NUMBER = "ArrayOfNumber"
+    ARRAY_OF_BOOLEAN = "ArrayOfBoolean"
+    NULL = "Null"
+
+    def __str__(self) -> str:
+        return self.value
 
+
+class AiAgentParameter:
     def __init__(
         self,
         name: str = None,
         description: str = None,
         send_to_model: bool = None,
+        policy: AiAgentParameterPolicy = AiAgentParameterPolicy.DEFAULT,
+        type: AiAgentParameterValueType = AiAgentParameterValueType.DEFAULT,
     ):
-        """
-        Initialize an agent parameter.
-
-        Args:
-            name: The parameter name. Cannot be null or empty.
-            description: A human-readable description. May be null or empty.
-            send_to_model: When False, the parameter is hidden from the model
-                          (it will not be included in prompts/echo messages).
-                          When True, the parameter is exposed to the model.
-                          If None (default), treated as exposed.
-        """
         self.name = name
         self.description: Optional[str] = description
         self.send_to_model: Optional[bool] = send_to_model
+        self.policy: AiAgentParameterPolicy = policy
+        self.type: AiAgentParameterValueType = type
 
     def to_json(self) -> Dict[str, Any]:
         return {
             "Name": self.name,
             "Description": self.description,
             "SendToModel": self.send_to_model,
+            "Policy": int(self.policy),
+            "Type": self.type.value,
         }
 
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentParameter:
+        # Server emits Policy as either int (1) or PascalCase ("ForbidModelGeneration").
+        policy_raw = json_dict.get("policy") if "policy" in json_dict else json_dict.get("Policy")
+        if policy_raw is None or policy_raw == 0 or policy_raw == "":
+            policy = AiAgentParameterPolicy.DEFAULT
+        elif isinstance(policy_raw, str):
+            snake = "".join("_" + c if i > 0 and c.isupper() else c for i, c in enumerate(policy_raw)).upper()
+            policy = AiAgentParameterPolicy[snake]
+        else:
+            policy = AiAgentParameterPolicy(policy_raw)
+
+        type_raw = json_dict.get("type") if "type" in json_dict else json_dict.get("Type")
+        type_ = AiAgentParameterValueType(type_raw) if type_raw else AiAgentParameterValueType.DEFAULT
+
         return cls(
             name=json_dict.get("name") or json_dict.get("Name"),
             description=json_dict.get("description") or json_dict.get("Description"),
             send_to_model=json_dict.get("sendToModel") if "sendToModel" in json_dict else json_dict.get("SendToModel"),
+            policy=policy,
+            type=type_,
         )
 
 
 class AiAgentToolQuery:
-    """
-    Represents a query tool that can be invoked by an AI agent.
-    The tool includes a name, description, query string, and parameter schema or sample object.
-    When invoked by the AI model, the query is expected to be executed by the server (database),
-    and its results provided back to the model.
-    """
-
+    # Database-side RQL the model can call. Results are sent back to the model.
     def __init__(
         self,
         name: str = None,
@@ -98,12 +119,8 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentToolQuery:
 
 
 class AiAgentToolAction:
-    """
-    Represents a tool action that can be invoked by an AI agent.
-    Includes metadata such as name, description, and optional parameters schema or sample.
-    Tool actions represent external functions whose results are provided by the user
-    """
-
+    # External function the model can call. Its result is supplied by the user
+    # (vs AiAgentToolQuery whose result comes from the database).
     def __init__(
         self,
         name: str = None,
@@ -137,11 +154,6 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentToolAction:
 
 
 class AiAgentPersistenceConfiguration:
-    """
-    Configuration for persisting chat history in RavenDB.
-    Defines where chat sessions should be stored and optionally how long they should be retained (expiration).
-    """
-
     def __init__(self, conversation_id_prefix: str = None, expires: int = None):
         self.conversation_id_prefix = conversation_id_prefix
         self.conversation_expiration_in_sec: Optional[int] = expires
@@ -163,10 +175,6 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentPersistenceConfiguration
 
 
 class AiAgentSummarizationByTokens:
-    """
-    Configuration settings for AI agent conversation summarization.
-    """
-
     DEFAULT_MAX_TOKENS_BEFORE_SUMMARIZATION = 32 * 1024
 
     def __init__(
@@ -208,10 +216,6 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentSummarizationByTokens:
 
 
 class AiAgentTruncateChat:
-    """
-    Configuration for truncating the AI chat history based on message count.
-    """
-
     DEFAULT_MESSAGES_LENGTH_BEFORE_TRUNCATE = 500
 
     def __init__(self, messages_length_before_truncate: int = None, messages_length_after_truncate: int = None):
@@ -241,10 +245,6 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentTruncateChat:
 
 
 class AiAgentHistoryConfiguration:
-    """
-    Defines the configuration for retention and expiration of AI agent chat history documents.
-    """
-
     def __init__(self, history_expiration_in_sec: int = None):
         self.history_expiration_in_sec: Optional[int] = history_expiration_in_sec
 
@@ -261,10 +261,6 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentHistoryConfiguration:
 
 
 class AiAgentChatTrimmingConfiguration:
-    """
-    Defines configuration options for reducing the size of the AI agent's chat history.
-    """
-
     def __init__(
         self,
         tokens_config: AiAgentSummarizationByTokens = None,
@@ -295,11 +291,6 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentChatTrimmingConfiguratio
 
 
 class AiAgentConfiguration:
-    """
-    Defines the configuration for an AI agent in RavenDB, including the system prompt,
-    tools (queries/actions), output schema, persistence settings, and connection string.
-    """
-
     def __init__(
         self,
         name: str = None,
@@ -363,7 +354,6 @@ def to_json(self) -> Dict[str, Any]:
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentConfiguration:
         instance = cls()
-        # Handle both camelCase and PascalCase for compatibility
         instance.identifier = json_dict.get("identifier") or json_dict.get("Identifier")
         instance.name = json_dict.get("name") or json_dict.get("Name")
         instance.connection_string_name = json_dict.get("connectionStringName") or json_dict.get("ConnectionStringName")

ravendb/tests/ai_agent_tests/test_ai_agent_configuration.py

@@ -0,0 +1,126 @@
+"""
+Unit tests for the 7.2.3 AI agent configuration additions:
+  * AiAgentParameter.policy / .type enums
+  * AiAgentToolSubAgent
+  * AiAgentConfiguration.sub_agents
+  * MissingAiAgentParameterException + dispatcher mapping
+  * AiMessagePromptFields.IMAGE
+  * AzureOpenAiSettings endpoint validation
+"""
+
+import unittest
+
+from ravendb import (
+    AiAgentConfiguration,
+    AiAgentParameter,
+    AiAgentParameterPolicy,
+    AiAgentParameterValueType,
+    AiAgentToolSubAgent,
+)
+from ravendb.documents.ai.content_part import AiMessagePromptFields
+from ravendb.documents.operations.ai.azure_open_ai_settings import AzureOpenAiSettings
+from ravendb.exceptions.exception_dispatcher import _EXCEPTION_MAP
+from ravendb.exceptions.raven_exceptions import MissingAiAgentParameterException
+
+
+class TestAiAgentParameterPolicyAndType(unittest.TestCase):
+    def test_defaults(self):
+        p = AiAgentParameter(name="userId")
+        self.assertEqual(AiAgentParameterPolicy.DEFAULT, p.policy)
+        self.assertEqual(AiAgentParameterValueType.DEFAULT, p.type)
+
+    def test_to_json_emits_policy_and_type(self):
+        p = AiAgentParameter(
+            name="userId",
+            description="Hidden",
+            send_to_model=False,
+            policy=AiAgentParameterPolicy.FORBID_MODEL_GENERATION,
+            type=AiAgentParameterValueType.STRING,
+        )
+        out = p.to_json()
+        self.assertEqual(1, out["Policy"])
+        self.assertEqual("String", out["Type"])
+
+    def test_from_json_round_trip(self):
+        out = {
+            "Name": "u",
+            "Description": "d",
+            "SendToModel": True,
+            "Policy": 1,
+            "Type": "ArrayOfNumber",
+        }
+        p = AiAgentParameter.from_json(out)
+        self.assertEqual(AiAgentParameterPolicy.FORBID_MODEL_GENERATION, p.policy)
+        self.assertEqual(AiAgentParameterValueType.ARRAY_OF_NUMBER, p.type)
+
+    def test_from_json_accepts_policy_as_string(self):
+        # The server returns Policy as the enum *name* string in GET responses
+        # (e.g. "Default", "ForbidModelGeneration"), even though Python emits
+        # the int on the way out. from_json must handle both forms.
+        p = AiAgentParameter.from_json({"Name": "u", "Policy": "ForbidModelGeneration"})
+        self.assertEqual(AiAgentParameterPolicy.FORBID_MODEL_GENERATION, p.policy)
+
+    def test_from_json_default_policy_as_string(self):
+        p = AiAgentParameter.from_json({"Name": "u", "Policy": "Default"})
+        self.assertEqual(AiAgentParameterPolicy.DEFAULT, p.policy)
+
+
+class TestAiAgentToolSubAgent(unittest.TestCase):
+    def test_round_trip(self):
+        sa = AiAgentToolSubAgent(identifier="benefits-agent", description="Handles benefits questions")
+        out = sa.to_json()
+        self.assertEqual("benefits-agent", out["Identifier"])
+        self.assertEqual("Handles benefits questions", out["Description"])
+        back = AiAgentToolSubAgent.from_json(out)
+        self.assertEqual(sa.identifier, back.identifier)
+        self.assertEqual(sa.description, back.description)
+
+
+class TestAiAgentConfigurationSubAgents(unittest.TestCase):
+    def test_to_json_includes_sub_agents(self):
+        cfg = AiAgentConfiguration(
+            name="main",
+            connection_string_name="cs",
+            system_prompt="be helpful",
+            sub_agents=[
+                AiAgentToolSubAgent("benefits", "benefits-related"),
+                AiAgentToolSubAgent("attendance", "PTO tracking"),
+            ],
+        )
+        out = cfg.to_json()
+        self.assertEqual(2, len(out["SubAgents"]))
+        self.assertEqual("benefits", out["SubAgents"][0]["Identifier"])
+
+
+class TestMissingAiAgentParameterException(unittest.TestCase):
+    def test_dispatcher_registers_short_name(self):
+        self.assertIn("MissingAiAgentParameterException", _EXCEPTION_MAP)
+        self.assertIs(MissingAiAgentParameterException, _EXCEPTION_MAP["MissingAiAgentParameterException"])
+
+    def test_is_raven_exception(self):
+        from ravendb.exceptions.raven_exceptions import RavenException
+
+        self.assertTrue(issubclass(MissingAiAgentParameterException, RavenException))
+
+
+class TestAiMessagePromptFieldsImage(unittest.TestCase):
+    def test_image_constant(self):
+        self.assertEqual("image", AiMessagePromptFields.IMAGE)
+
+
+class TestAzureOpenAiSettingsEndpointValidation(unittest.TestCase):
+    def test_none_endpoint_rejected(self):
+        with self.assertRaises(ValueError):
+            AzureOpenAiSettings(api_key="k", endpoint=None, model="m", deployment_name="d")
+
+    def test_blank_endpoint_rejected(self):
+        with self.assertRaises(ValueError):
+            AzureOpenAiSettings(api_key="k", endpoint="   ", model="m", deployment_name="d")
+
+    def test_valid_endpoint_accepted(self):
+        s = AzureOpenAiSettings(api_key="k", endpoint="https://x.openai.azure.com", model="m", deployment_name="d")
+        self.assertEqual("https://x.openai.azure.com", s.endpoint)
+
+
+if __name__ == "__main__":
+    unittest.main()