langroid 0.23.3__py3-none-any.whl → 0.24.1__py3-none-any.whl

langroid/agent/chat_document.py

@@ -217,9 +217,14 @@ class ChatDocument(Document):
         if self.function_call is not None:
             tool_type = "FUNC"
             tool = self.function_call.name
-        elif (json_tools := self.get_tool_names()) != []:
-            tool_type = "TOOL"
-            tool = json_tools[0]
+        else:
+            try:
+                json_tools = self.get_tool_names()
+            except Exception:
+                json_tools = []
+            if json_tools != []:
+                tool_type = "TOOL"
+                tool = json_tools[0]
         recipient = self.metadata.recipient
         content = self.content
         sender_entity = self.metadata.sender

langroid/agent/openai_assistant.py

@@ -203,7 +203,7 @@ class OpenAIAssistant(ChatAgent):
             self.set_system_message(sys_msg.content)
         if not self.config.use_functions_api:
             return
-        functions, _, _, _ = self._function_args()
+        functions, _, _, _, _ = self._function_args()
         if functions is None:
             return
         # add the functions to the assistant:

langroid/agent/special/sql/sql_chat_agent.py

@@ -359,13 +359,27 @@ class SQLChatAgent(ChatAgent):
         # This is likelier to succeed since this agent has no "baggage" of
         # prior conversation, other than the system msg, and special
         # "Intent-interpretation" instructions.
-        response = self.helper_agent.llm_response(message)
-        tools = self.try_get_tool_messages(response)
-        if tools:
-            return response
+        if self._json_schema_available():
+            AnyTool = self._get_any_tool_message(optional=False)
+            self.set_output_format(
+                AnyTool,
+                force_tools=True,
+                use=True,
+                handle=True,
+                instructions=True,
+            )
+            recovery_message = self._strict_recovery_instructions(
+                AnyTool, optional=False
+            )
+            return self.llm_response(recovery_message)
         else:
-            # fall back on the clarification message
-            return self._clarifying_message()
+            response = self.helper_agent.llm_response(message)
+            tools = self.try_get_tool_messages(response)
+            if tools:
+                return response
+            else:
+                # fall back on the clarification message
+                return self._clarifying_message()
 
     def retry_query(self, e: Exception, query: str) -> str:
         """

langroid/agent/task.py

@@ -16,6 +16,7 @@ from typing import (
     Dict,
     List,
     Optional,
+    Self,
     Tuple,
     Type,
     TypeVar,
@@ -598,7 +599,7 @@ class Task:
         for t in self.sub_tasks:
             t.reset_all_sub_tasks()
 
-    def __getitem__(self, return_type: type) -> Task:
+    def __getitem__(self, return_type: type) -> Self:
         """Returns a (shallow) copy of `self` with a default return type."""
         clone = copy.copy(self)
         clone.default_return_type = return_type
@@ -732,8 +733,37 @@ class Task:
         if return_type is None:
             return_type = self.default_return_type
 
+        # If possible, take a final strict decoding step
+        # when the output does not match `return_type`
         if return_type is not None and return_type != ChatDocument:
-            return self.agent.from_ChatDocument(final_result, return_type)
+            parsed_result = self.agent.from_ChatDocument(final_result, return_type)
+
+            if (
+                parsed_result is None
+                and isinstance(self.agent, ChatAgent)
+                and self.agent._json_schema_available()
+            ):
+                strict_agent = self.agent[return_type]
+                output_args = strict_agent._function_args()[-1]
+                if output_args is not None:
+                    schema = output_args.function.parameters
+                    strict_result = strict_agent.llm_response(
+                        f"""
+                        A response adhering to the following JSON schema was expected:
+                        {schema}
+
+                        Please resubmit with the correct schema. 
+                        """
+                    )
+
+                    if strict_result is not None:
+                        return cast(
+                            Optional[T],
+                            strict_agent.from_ChatDocument(strict_result, return_type),
+                        )
+
+            return parsed_result
+
         return final_result
 
     @overload
@@ -895,8 +925,37 @@ class Task:
         if return_type is None:
             return_type = self.default_return_type
 
+        # If possible, take a final strict decoding step
+        # when the output does not match `return_type`
         if return_type is not None and return_type != ChatDocument:
-            return self.agent.from_ChatDocument(final_result, return_type)
+            parsed_result = self.agent.from_ChatDocument(final_result, return_type)
+
+            if (
+                parsed_result is None
+                and isinstance(self.agent, ChatAgent)
+                and self.agent._json_schema_available()
+            ):
+                strict_agent = self.agent[return_type]
+                output_args = strict_agent._function_args()[-1]
+                if output_args is not None:
+                    schema = output_args.function.parameters
+                    strict_result = await strict_agent.llm_response_async(
+                        f"""
+                        A response adhering to the following JSON schema was expected:
+                        {schema}
+
+                        Please resubmit with the correct schema. 
+                        """
+                    )
+
+                    if strict_result is not None:
+                        return cast(
+                            Optional[T],
+                            strict_agent.from_ChatDocument(strict_result, return_type),
+                        )
+
+            return parsed_result
+
         return final_result
 
     def _pre_run_loop(

langroid/agent/tool_message.py

@@ -6,11 +6,12 @@ an agent. The messages could represent, for example:
 - request to run a method of the agent
 """
 
+import copy
 import json
 import textwrap
 from abc import ABC
 from random import choice
-from typing import Any, Dict, List, Tuple, Type
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar
 
 from docstring_parser import parse
 
@@ -22,6 +23,55 @@ from langroid.utils.pydantic_utils import (
 )
 from langroid.utils.types import is_instance_of
 
+K = TypeVar("K")
+
+
+def remove_if_exists(k: K, d: dict[K, Any]) -> None:
+    """Removes key `k` from `d` if present."""
+    if k in d:
+        d.pop(k)
+
+
+def format_schema_for_strict(schema: Any) -> None:
+    """
+    Recursively set additionalProperties to False and replace
+    oneOf and allOf with anyOf, required for OpenAI structured outputs.
+    Additionally, remove all defaults and set all fields to required.
+    This may not be equivalent to the original schema.
+    """
+    if isinstance(schema, dict):
+        if "type" in schema and schema["type"] == "object":
+            schema["additionalProperties"] = False
+
+            if "properties" in schema:
+                properties = schema["properties"]
+                all_properties = list(properties.keys())
+                for k, v in properties.items():
+                    if "default" in v:
+                        if k == "request":
+                            v["enum"] = [v["default"]]
+
+                        v.pop("default")
+                schema["required"] = all_properties
+            else:
+                schema["properties"] = {}
+                schema["required"] = []
+
+        anyOf = (
+            schema.get("oneOf", []) + schema.get("allOf", []) + schema.get("anyOf", [])
+        )
+        if "allOf" in schema or "oneOf" in schema or "anyOf" in schema:
+            schema["anyOf"] = anyOf
+
+        remove_if_exists("allOf", schema)
+        remove_if_exists("oneOf", schema)
+
+        for v in schema.values():
+            format_schema_for_strict(v)
+    elif isinstance(schema, list):
+        for v in schema:
+            format_schema_for_strict(v)
+
 
 class ToolMessage(ABC, BaseModel):
     """
@@ -42,6 +92,9 @@ class ToolMessage(ABC, BaseModel):
     purpose: str
     id: str = ""  # placeholder for OpenAI-API tool_call_id
 
+    # If enabled, forces strict adherence to schema.
+    # Currently only supported by OpenAI LLMs. When unset, enables if supported.
+    _strict: Optional[bool] = None
     _allow_llm_use: bool = True  # allow an LLM to use (i.e. generate) this tool?
 
     # Optional param to limit number of result tokens to retain in msg history.
@@ -239,7 +292,7 @@ class ToolMessage(ABC, BaseModel):
             LLMFunctionSpec: the schema as an LLMFunctionSpec
 
         """
-        schema = cls.schema()
+        schema = copy.deepcopy(cls.schema())
         docstring = parse(cls.__doc__ or "")
         parameters = {
             k: v for k, v in schema.items() if k not in ("title", "description")
@@ -268,6 +321,13 @@ class ToolMessage(ABC, BaseModel):
         if request:
             parameters["required"].append("request")
 
+            # If request is present it must match the default value
+            # Similar to defining request as a literal type
+            parameters["request"] = {
+                "enum": [cls.default_value("request")],
+                "type": "string",
+            }
+
         if "description" not in schema:
             if docstring.short_description:
                 schema["description"] = docstring.short_description
@@ -277,6 +337,26 @@ class ToolMessage(ABC, BaseModel):
                     f"the required parameters with correct types"
                 )
 
+        # Handle nested ToolMessage fields
+        if "definitions" in parameters:
+            for v in parameters["definitions"].values():
+                if "exclude" in v:
+                    v.pop("exclude")
+
+                    remove_if_exists("purpose", v["properties"])
+                    remove_if_exists("id", v["properties"])
+                    if (
+                        "request" in v["properties"]
+                        and "default" in v["properties"]["request"]
+                    ):
+                        if "required" not in v:
+                            v["required"] = []
+                        v["required"].append("request")
+                        v["properties"]["request"] = {
+                            "type": "string",
+                            "enum": [v["properties"]["request"]["default"]],
+                        }
+
         parameters.pop("exclude")
         _recursive_purge_dict_key(parameters, "title")
         _recursive_purge_dict_key(parameters, "additionalProperties")

langroid/agent/tools/orchestration.py

@@ -97,7 +97,7 @@ class ResultTool(ToolMessage):
         validate_assignment = True
         # do not include these fields in the generated schema
         # since we don't require the LLM to specify them
-        schema_extra = {"exclude": {"purpose", "id"}}
+        schema_extra = {"exclude": {"purpose", "id", "strict"}}
 
     def handle(self) -> AgentDoneTool:
         return AgentDoneTool(tools=[self])
@@ -134,7 +134,7 @@ class FinalResultTool(ToolMessage):
         validate_assignment = True
         # do not include these fields in the generated schema
         # since we don't require the LLM to specify them
-        schema_extra = {"exclude": {"purpose", "id"}}
+        schema_extra = {"exclude": {"purpose", "id", "strict"}}
 
 
 class PassTool(ToolMessage):

langroid/agent/xml_tool_message.py

@@ -1,3 +1,4 @@
+import re
 from collections.abc import Mapping
 from typing import Any, Dict, List, Optional, get_args, get_origin
 
@@ -323,45 +324,59 @@ class XMLToolMessage(ToolMessage):
     @classmethod
     def find_candidates(cls, text: str) -> List[str]:
         """
-        Find and extract all potential XML tool messages from the given text.
-
-        This method searches for XML-like structures in the input text that match
-        the expected format of the tool message. It looks for opening and closing
-        tags that correspond to the root element defined in the XMLToolMessage class,
-        which is by default <tool>.
+        Finds XML-like tool message candidates in text, with relaxed opening tag rules.
 
         Args:
-            text (str): The input text to search for XML tool messages.
+            text: Input text to search for XML structures.
 
         Returns:
-            List[str]: A list of strings, each representing a potential XML tool
-                       message.
-                       These candidates include both the opening and
-                       closing tags, so that they are individually parseable.
-
-        Note:
-            This method ensures that all candidates are valid and parseable by
-            inserting a closing tag if it's missing for the last candidate.
+            List of XML strings. For fragments missing the root opening tag but having
+            valid XML structure and root closing tag, prepends the root opening tag.
+
+        Example:
+            With root_tag="tool", given:
+            "Hello <field1>data</field1> </tool>"
+            Returns: ["<tool><field1>data</field1></tool>"]
         """
+
         root_tag = cls.Config.root_element
         opening_tag = f"<{root_tag}>"
         closing_tag = f"</{root_tag}>"
 
         candidates = []
-        start = 0
+        pos = 0
         while True:
-            start = text.find(opening_tag, start)
-            if start == -1:
-                break
-            end = text.find(closing_tag, start)
-            if end == -1:
-                # For the last candidate, insert the closing tag if it's missing
-                candidate = text[start:]
-                if not candidate.strip().endswith(closing_tag):
-                    candidate += closing_tag
-                candidates.append(candidate)
+            # Look for either proper opening tag or closing tag
+            start_normal = text.find(opening_tag, pos)
+            end = text.find(closing_tag, pos)
+
+            if start_normal == -1 and end == -1:
                 break
-            candidates.append(text[start : end + len(closing_tag)])
-            start = end + len(closing_tag)
+
+            if start_normal != -1:
+                # Handle normal case (has opening tag)
+                end = text.find(closing_tag, start_normal)
+                if end != -1:
+                    candidates.append(text[start_normal : end + len(closing_tag)])
+                    pos = max(end + len(closing_tag), start_normal + 1)
+                    continue
+                elif start_normal == text.rfind(opening_tag):
+                    # last fragment - ok to miss closing tag
+                    candidates.append(text[start_normal:] + closing_tag)
+                    return candidates
+                else:
+                    pos = start_normal + 1
+                    continue
+
+            if end != -1:
+                # Look backwards for first XML tag
+                text_before = text[pos:end]
+                first_tag_match = re.search(r"<\w+>", text_before)
+                if first_tag_match:
+                    start = pos + first_tag_match.start()
+                    candidates.append(
+                        opening_tag + text[start : end + len(closing_tag)]
+                    )
+                pos = end + len(closing_tag)
 
         return candidates

langroid/language_models/azure_openai.py

@@ -8,6 +8,13 @@ from langroid.language_models.openai_gpt import (
     OpenAIGPTConfig,
 )
 
+azureStructuredOutputList = [
+    "2024-08-06",
+    "2024-11-20",
+]
+
+azureStructuredOutputAPIMin = "2024-08-01-preview"
+
 
 class AzureConfig(OpenAIGPTConfig):
     """
@@ -96,6 +103,11 @@ class AzureGPT(OpenAIGPT):
         # when you deployed a model
         self.set_chat_model()
 
+        self.supports_json_schema = (
+            self.config.api_version >= azureStructuredOutputAPIMin
+            and self.config.model_version in azureStructuredOutputList
+        )
+
         self.client = AzureOpenAI(
             api_key=self.config.api_key,
             azure_endpoint=self.config.api_base,
@@ -136,12 +148,13 @@ class AzureGPT(OpenAIGPT):
         If the version is not set, it raises a ValueError indicating
         that the model version needs to be specified in the ``.env``
         file.  It sets `OpenAIChatMode.GPT4o` if the version is
-        '2024-05-13', `OpenAIChatModel.GPT4_TURBO` if the version is
-        '1106-Preview', otherwise, it defaults to setting
-        `OpenAIChatModel.GPT4`.
+        one of those listed below, and
+        `OpenAIChatModel.GPT4_TURBO` if
+        the version is '1106-Preview', otherwise, it defaults to
+        setting `OpenAIChatModel.GPT4`.
         """
         VERSION_1106_PREVIEW = "1106-Preview"
-        VERSION_GPT4o = "2024-05-13"
+        VERSIONS_GPT4o = ["2024-05-13", "2024-08-06", "2024-11-20"]
 
         if self.config.model_version == "":
             raise ValueError(
@@ -149,7 +162,7 @@ class AzureGPT(OpenAIGPT):
                 "Please set it to the chat model version used in your deployment."
             )
 
-        if self.config.model_version == VERSION_GPT4o:
+        if self.config.model_version in VERSIONS_GPT4o:
             self.config.chat_model = OpenAIChatModel.GPT4o
         elif self.config.model_version == VERSION_1106_PREVIEW:
             self.config.chat_model = OpenAIChatModel.GPT4_TURBO

langroid/language_models/base.py

@@ -156,9 +156,29 @@ class OpenAIToolCall(BaseModel):
 
 class OpenAIToolSpec(BaseModel):
     type: ToolTypes
+    strict: Optional[bool] = None
     function: LLMFunctionSpec
 
 
+class OpenAIJsonSchemaSpec(BaseModel):
+    strict: Optional[bool] = None
+    function: LLMFunctionSpec
+
+    def to_dict(self) -> Dict[str, Any]:
+        json_schema: Dict[str, Any] = {
+            "name": self.function.name,
+            "description": self.function.description,
+            "schema": self.function.parameters,
+        }
+        if self.strict is not None:
+            json_schema["strict"] = self.strict
+
+        return {
+            "type": "json_schema",
+            "json_schema": json_schema,
+        }
+
+
 class LLMTokenUsage(BaseModel):
     """
     Usage of tokens by an LLM.
@@ -512,6 +532,7 @@ class LanguageModel(ABC):
         tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
         functions: Optional[List[LLMFunctionSpec]] = None,
         function_call: str | Dict[str, str] = "auto",
+        response_format: Optional[OpenAIJsonSchemaSpec] = None,
     ) -> LLMResponse:
         """
         Get chat-completion response from LLM.
@@ -538,6 +559,7 @@ class LanguageModel(ABC):
         tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
         functions: Optional[List[LLMFunctionSpec]] = None,
         function_call: str | Dict[str, str] = "auto",
+        response_format: Optional[OpenAIJsonSchemaSpec] = None,
     ) -> LLMResponse:
         """Async version of `chat`. See `chat` for details."""
         pass

langroid/language_models/mock_lm.py

@@ -7,6 +7,7 @@ from langroid.language_models import LLMResponse
 from langroid.language_models.base import (
     LanguageModel,
     LLMConfig,
+    OpenAIJsonSchemaSpec,
     OpenAIToolSpec,
     ToolChoiceTypes,
 )
@@ -80,6 +81,7 @@ class MockLM(LanguageModel):
         tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
         functions: Optional[List[lm.LLMFunctionSpec]] = None,
         function_call: str | Dict[str, str] = "auto",
+        response_format: Optional[OpenAIJsonSchemaSpec] = None,
     ) -> lm.LLMResponse:
         """
         Mock chat function for testing
@@ -95,6 +97,7 @@ class MockLM(LanguageModel):
         tool_choice: ToolChoiceTypes | Dict[str, str | Dict[str, str]] = "auto",
         functions: Optional[List[lm.LLMFunctionSpec]] = None,
         function_call: str | Dict[str, str] = "auto",
+        response_format: Optional[OpenAIJsonSchemaSpec] = None,
     ) -> lm.LLMResponse:
         """
         Mock chat function for testing