RDBC-1059 Add AiAgentToolSubAgent and AiAgentConfiguration.sub_agents

Author: poissoncorp

ravendb/__init__.py

@@ -95,6 +95,7 @@
     AiAgentParameterValueType,
     AiAgentToolAction,
     AiAgentToolQuery,
+    AiAgentToolSubAgent,
     AiAgentPersistenceConfiguration,
     AiAgentChatTrimmingConfiguration,
     AiAgentSummarizationByTokens,

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

@@ -12,6 +12,7 @@
     AiAgentTruncateChat,
     AiAgentHistoryConfiguration,
 )
+from .ai_agent_tool_sub_agent import AiAgentToolSubAgent
 
 from .add_or_update_ai_agent_operation import (
     AddOrUpdateAiAgentOperation,
@@ -45,6 +46,7 @@
     "AiAgentToolAction",
     "AiAgentToolQuery",
     "AiAgentToolQueryOptions",
+    "AiAgentToolSubAgent",
     "AiAgentPersistenceConfiguration",
     "AiAgentChatTrimmingConfiguration",
     "AiAgentSummarizationByTokens",

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

@@ -306,7 +306,10 @@ def __init__(
         chat_trimming: AiAgentChatTrimmingConfiguration = None,
         max_model_iterations_per_call: int = None,
         disabled: bool = False,
+        sub_agents: List["AiAgentToolSubAgent"] = None,
     ):
+        from ravendb.documents.operations.ai.agents.ai_agent_tool_sub_agent import AiAgentToolSubAgent
+
         self.name = name
         self.connection_string_name = connection_string_name
         self.system_prompt = system_prompt
@@ -320,10 +323,10 @@ def __init__(
         self.chat_trimming: Optional[AiAgentChatTrimmingConfiguration] = chat_trimming
         self.max_model_iterations_per_call: Optional[int] = max_model_iterations_per_call
         self.disabled: bool = disabled
+        self.sub_agents: List[AiAgentToolSubAgent] = sub_agents or []
 
     @staticmethod
     def _normalize_parameters(parameters: List[Union[str, AiAgentParameter]]) -> List[AiAgentParameter]:
-        """Convert a list of strings or AiAgentParameter objects to a list of AiAgentParameter objects."""
         if not parameters:
             return []
         result = []
@@ -344,6 +347,7 @@ def to_json(self) -> Dict[str, Any]:
             "OutputSchema": self.output_schema,
             "Queries": [q.to_json() for q in self.queries],
             "Actions": [a.to_json() for a in self.actions],
+            "SubAgents": [s.to_json() for s in self.sub_agents],
             "Persistence": self.persistence.to_json() if self.persistence else None,
             "Parameters": [p.to_json() for p in self.parameters],
             "ChatTrimming": self.chat_trimming.to_json() if self.chat_trimming else None,
@@ -369,6 +373,12 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentConfiguration:
         if actions_data:
             instance.actions = [AiAgentToolAction.from_json(a) for a in actions_data]
 
+        from ravendb.documents.operations.ai.agents.ai_agent_tool_sub_agent import AiAgentToolSubAgent
+
+        sub_agents_data = json_dict.get("subAgents") or json_dict.get("SubAgents")
+        if sub_agents_data:
+            instance.sub_agents = [AiAgentToolSubAgent.from_json(s) for s in sub_agents_data]
+
         persistence_data = json_dict.get("persistence") or json_dict.get("Persistence")
         if persistence_data:
             instance.persistence = AiAgentPersistenceConfiguration.from_json(persistence_data)

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

@@ -0,0 +1,23 @@
+from __future__ import annotations
+from typing import Any, Dict, Optional
+
+
+class AiAgentToolSubAgent:
+    """Server-side sub-agent the model can call from within a parent agent run."""
+
+    def __init__(self, identifier: Optional[str] = None, description: Optional[str] = None):
+        self.identifier = identifier
+        self.description = description
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "Identifier": self.identifier,
+            "Description": self.description,
+        }
+
+    @classmethod
+    def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentToolSubAgent:
+        return cls(
+            identifier=json_dict.get("identifier") or json_dict.get("Identifier"),
+            description=json_dict.get("description") or json_dict.get("Description"),
+        )

ravendb/tests/ai_agent_tests/test_ai_agent_config_extensions_integration.py

@@ -0,0 +1,137 @@
+"""
+Integration tests against a live RavenDB 7.2.x server for the AI agent
+configuration fields added in 7.2.3:
+  * AiAgentConfiguration.sub_agents (+ AiAgentToolSubAgent)
+  * AiAgentParameter.policy (AiAgentParameterPolicy)
+  * AiAgentParameter.type (AiAgentParameterValueType)
+
+The license guard mirrors the existing AI agent tests in this directory.
+"""
+
+import os
+import unittest
+
+from ravendb import (
+    AiAgentConfiguration,
+    AiAgentParameter,
+    AiAgentParameterPolicy,
+    AiAgentParameterValueType,
+    AiAgentToolSubAgent,
+)
+from ravendb.documents.operations.ai import AiConnectionString, AiModelType
+from ravendb.documents.operations.ai.agents import (
+    AddOrUpdateAiAgentOperation,
+    DeleteAiAgentOperation,
+)
+from ravendb.documents.operations.ai.open_ai_settings import OpenAiSettings
+from ravendb.documents.operations.connection_string.put_connection_string_operation import PutConnectionStringOperation
+from ravendb.tests.test_base import TestBase
+
+
+@unittest.skipIf(os.environ.get("RAVENDB_LICENSE") is None, "Insufficient license permissions. Skipping on CI/CD.")
+class TestAiAgentConfigExtensionsIntegration(TestBase):
+    CONNECTION_STRING_NAME = "test-ai-agent-cs-ext"
+
+    def setUp(self):
+        super().setUp()
+        ai_connection_string = AiConnectionString(
+            name=self.CONNECTION_STRING_NAME,
+            identifier=self.CONNECTION_STRING_NAME,
+            model_type=AiModelType.CHAT,
+            openai_settings=OpenAiSettings(
+                api_key="dummy-api-key",
+                endpoint="https://api.openai.com/v1",
+                model="gpt-4",
+            ),
+        )
+        self.store.maintenance.send(PutConnectionStringOperation(ai_connection_string))
+        self._created_agent_ids = []
+
+    def tearDown(self):
+        for agent_id in self._created_agent_ids:
+            try:
+                self.store.maintenance.send(DeleteAiAgentOperation(agent_id))
+            except Exception:
+                pass
+        super().tearDown()
+
+    # ---- sub_agents ----
+
+    def test_sub_agents_round_trip_through_get_agent(self):
+        agent = AiAgentConfiguration(
+            name="ParentAgent",
+            identifier="test-sub-agent-parent",
+            connection_string_name=self.CONNECTION_STRING_NAME,
+            system_prompt="Dispatch to sub-agents as needed.",
+            sample_object='{"answer": "..."}',
+            sub_agents=[
+                AiAgentToolSubAgent(identifier="benefits-agent", description="Handles benefit questions"),
+                AiAgentToolSubAgent(identifier="attendance-agent", description="Tracks PTO and attendance"),
+            ],
+        )
+        result = self.store.ai.add_or_update_agent(agent)
+        self._created_agent_ids.append(result.identifier)
+
+        fetched = self.store.ai.get_agents(result.identifier).ai_agents[0]
+        identifiers = sorted(s.identifier for s in fetched.sub_agents)
+        self.assertEqual(["attendance-agent", "benefits-agent"], identifiers)
+
+        descriptions = {s.identifier: s.description for s in fetched.sub_agents}
+        self.assertEqual("Handles benefit questions", descriptions["benefits-agent"])
+        self.assertEqual("Tracks PTO and attendance", descriptions["attendance-agent"])
+
+    # ---- AiAgentParameter.policy ----
+
+    def test_parameter_policy_forbid_model_generation_round_trips(self):
+        agent = AiAgentConfiguration(
+            name="ParamPolicyAgent",
+            identifier="test-param-policy",
+            connection_string_name=self.CONNECTION_STRING_NAME,
+            system_prompt="Test parameter policy.",
+            sample_object='{"answer": "..."}',
+            parameters=[
+                AiAgentParameter(
+                    name="user_id",
+                    description="Hidden user id",
+                    send_to_model=False,
+                    policy=AiAgentParameterPolicy.FORBID_MODEL_GENERATION,
+                ),
+                AiAgentParameter(name="country", description="The country to filter by."),
+            ],
+        )
+        result = self.store.ai.add_or_update_agent(agent)
+        self._created_agent_ids.append(result.identifier)
+
+        fetched = self.store.ai.get_agents(result.identifier).ai_agents[0]
+        by_name = {p.name: p for p in fetched.parameters}
+        self.assertEqual(AiAgentParameterPolicy.FORBID_MODEL_GENERATION, by_name["user_id"].policy)
+        # Default-valued parameter still comes back with the default policy.
+        self.assertEqual(AiAgentParameterPolicy.DEFAULT, by_name["country"].policy)
+
+    # ---- AiAgentParameter.type ----
+
+    def test_parameter_value_type_round_trips(self):
+        agent = AiAgentConfiguration(
+            name="ParamTypeAgent",
+            identifier="test-param-type",
+            connection_string_name=self.CONNECTION_STRING_NAME,
+            system_prompt="Test parameter types.",
+            sample_object='{"answer": "..."}',
+            parameters=[
+                AiAgentParameter(name="email", description="Email", type=AiAgentParameterValueType.STRING),
+                AiAgentParameter(name="age", description="Age", type=AiAgentParameterValueType.NUMBER),
+                AiAgentParameter(name="tags", description="Tags", type=AiAgentParameterValueType.ARRAY_OF_STRING),
+            ],
+        )
+        result = self.store.ai.add_or_update_agent(agent)
+        self._created_agent_ids.append(result.identifier)
+
+        fetched = self.store.ai.get_agents(result.identifier).ai_agents[0]
+        by_name = {p.name: p for p in fetched.parameters}
+        self.assertEqual(AiAgentParameterValueType.STRING, by_name["email"].type)
+        self.assertEqual(AiAgentParameterValueType.NUMBER, by_name["age"].type)
+        self.assertEqual(AiAgentParameterValueType.ARRAY_OF_STRING, by_name["tags"].type)
+
+
+if __name__ == "__main__":
+    unittest.main()