waldiez 0.5.3__py3-none-any.whl → 0.5.5__py3-none-any.whl

waldiez/_version.py

@@ -5,4 +5,4 @@
 This file is automatically generated by Hatchling.
 Do not edit this file directly.
 """
-__version__ = VERSION = "0.5.3"
+__version__ = VERSION = "0.5.5"

waldiez/cli.py

@@ -92,21 +92,8 @@ def run(
     structured: bool = typer.Option(  # noqa: B008
         False,
         help=(
-            "If set, running the flow will use structured io stream instead of the default 'input/print' "
-        ),
-    ),
-    threaded: bool = typer.Option(  # noqa: B008
-        False,
-        help=(
-            "If set, the flow will be run in a separate thread. "
-            "This is useful for running flows that require user input or print output."
-        ),
-    ),
-    patch_io: bool = typer.Option(  # noqa: B008
-        False,
-        help=(
-            "If set, the flow will patch ag2's IOStream to safe print and input methods. "
-            "This is useful for running flows that require user input or print output."
+            "If set, the output will be structured as a directory with "
+            "the flow file and any additional generated files in it."
         ),
     ),
     force: bool = typer.Option(  # noqa: B008
@@ -136,8 +123,7 @@ def run(
             runner.a_run,
             output_path,
             uploads_root,
-            structured,
-            not patch_io,  # skip_patch_io
+            structured,  # structured_io
             False,  # skip_mmd
             False,  # skip_timeline
         )
@@ -146,8 +132,6 @@ def run(
             output_path=output_path,
             uploads_root=uploads_root,
             structured_io=structured,
-            threaded=threaded,
-            skip_patch_io=not patch_io,
             skip_mmd=False,
             skip_timeline=False,
         )
@@ -184,13 +168,6 @@ def convert(
         False,
         help="Override the output file if it already exists.",
     ),
-    patch_io: bool = typer.Option(  # noqa: B008
-        False,
-        help=(
-            "If set, the exported script will inlclude a snippet to patch ag2's IOStream "
-            "to safe print and input methods. "
-        ),
-    ),
     debug: bool = typer.Option(
         False,
         "--debug",
@@ -219,7 +196,6 @@ def convert(
     exporter.export(
         output,
         force=force,
-        skip_patch_io=not patch_io,
         debug=debug,
     )
     generated = str(output).replace(os.getcwd(), ".")

waldiez/exporter.py

@@ -63,7 +63,6 @@ class WaldiezExporter:
         path: str | Path,
         structured_io: bool = False,
         uploads_root: Path | None = None,
-        skip_patch_io: bool = True,
         force: bool = False,
         debug: bool = False,
     ) -> None:
@@ -78,8 +77,6 @@ class WaldiezExporter:
             by default False.
         uploads_root : str | Path | None, (optional)
             The uploads root, to get user-uploaded files, by default None.
-        skip_patch_io : bool, (optional)
-            Whether to skip patching I/O, by default True.
         force : bool, (optional)
             Override the output file if it already exists, by default False.
         debug : bool, (optional)
@@ -112,7 +109,6 @@ class WaldiezExporter:
                 path,
                 structured_io=structured_io,
                 uploads_root=uploads_root,
-                skip_patch_io=skip_patch_io,
                 debug=debug,
             )
         elif extension == ".ipynb":
@@ -120,7 +116,6 @@ class WaldiezExporter:
                 path,
                 structured_io=structured_io,
                 uploads_root=uploads_root,
-                skip_patch_io=skip_patch_io,
                 debug=debug,
             )
         else:
@@ -131,7 +126,6 @@ class WaldiezExporter:
         path: str | Path,
         structured_io: bool = False,
         uploads_root: Path | None = None,
-        skip_patch_io: bool = True,
         debug: bool = False,
     ) -> None:
         """Export flow to jupyter notebook.
@@ -145,8 +139,6 @@ class WaldiezExporter:
             by default False.
         uploads_root : Path | None, optional
             The uploads root, to get user-uploaded files, by default None.
-        skip_patch_io : bool, optional
-            Whether to skip patching I/O, by default True.
         debug : bool, optional
             Whether to enable debug mode, by default False.
 
@@ -165,7 +157,6 @@ class WaldiezExporter:
             uploads_root=uploads_root,
             structured_io=structured_io,
             for_notebook=True,
-            skip_patch_io=skip_patch_io,
             debug=debug,
         )
         self.flow_extras = exporter.extras
@@ -202,7 +193,6 @@ class WaldiezExporter:
         path: str | Path,
         structured_io: bool = False,
         uploads_root: Path | None = None,
-        skip_patch_io: bool = True,
         debug: bool = False,
     ) -> None:
         """Export waldiez flow to a python script.
@@ -216,8 +206,6 @@ class WaldiezExporter:
             by default False.
         uploads_root : Path | None, optional
             The uploads root, to get user-uploaded files, by default None.
-        skip_patch_io : bool, optional
-            Whether to skip patching I/O, by default True.
         debug : bool, optional
             Whether to enable debug mode, by default False.
 
@@ -234,7 +222,6 @@ class WaldiezExporter:
             for_notebook=False,
             uploads_root=uploads_root,
             structured_io=structured_io,
-            skip_patch_io=skip_patch_io,
             debug=debug,
         )
         self.flow_extras = exporter.extras

waldiez/exporting/agent/exporter.py

@@ -23,6 +23,7 @@ from ..core import (
 )
 from .code_execution import CodeExecutionProcessor
 from .extras.captain_agent_extras import CaptainAgentProcessor
+from .extras.doc_agent_extras import DocAgentProcessor
 from .extras.group_manager_agent_extas import GroupManagerProcessor
 from .extras.group_member_extras import GroupMemberAgentProcessor
 from .extras.rag_user_proxy_agent_extras import RagUserProxyAgentProcessor
@@ -125,6 +126,8 @@ class AgentExporter(Exporter[StandardExtras]):
             return self.create_reasoning_extras()
         if self.agent.is_rag_user:
             return self._create_rag_user_extras()
+        if self.agent.is_doc_agent:
+            return self._create_doc_agent_extras()
         return self._create_standard_extras()
 
     def _create_standard_extras(self) -> StandardExtras:
@@ -227,6 +230,41 @@ class AgentExporter(Exporter[StandardExtras]):
             extras.prepend_before_agent(rag_result.before_agent)
         return extras
 
+    def _create_doc_agent_extras(self) -> StandardExtras:
+        """Create doc agent extras."""
+        # Start with standard extras
+        extras = self._create_standard_extras()
+
+        doc_processor = DocAgentProcessor(
+            agent=self.agent,
+            agent_names=self.agent_names,
+            model_names=self.model_names,
+            all_models=self.models,
+            output_dir=self.output_dir,
+            cache_seed=self.cache_seed,
+            serializer=(
+                self.context.get_serializer()
+                if self.context.serializer
+                else None
+            ),
+        )
+        doc_extras = doc_processor.process(
+            system_message_config=extras.system_message_config,
+            termination_config=extras.termination_config,
+            code_execution_config=extras.code_execution_config,
+        )
+        # Add doc agent specific arguments
+        for arg in doc_extras.extra_args:
+            extras.add_arg(arg)
+        # Add doc agent specific imports
+        extras.add_imports(doc_extras.extra_imports)
+        if doc_extras.before_agent:
+            extras.prepend_before_agent(doc_extras.before_agent)
+        if doc_extras.after_agent:
+            extras.append_after_agent(doc_extras.after_agent)
+
+        return extras
+
     def _create_captain_extras(self) -> StandardExtras:
         """Create captain agent extras."""
         # Start with standard extras

waldiez/exporting/agent/extras/__init__.py

@@ -3,12 +3,14 @@
 """Agents exporter extras."""
 
 from .captain_agent_extras import CaptainAgentProcessor
+from .doc_agent_extras import DocAgentProcessor
 from .group_manager_agent_extas import GroupManagerProcessor
 from .group_member_extras import GroupMemberAgentProcessor
 from .reasoning_agent_extras import ReasoningAgentProcessor
 
 __all__ = [
     "CaptainAgentProcessor",
+    "DocAgentProcessor",
     "GroupManagerProcessor",
     "GroupMemberAgentProcessor",
     "ReasoningAgentProcessor",

waldiez/exporting/agent/extras/doc_agent_extras.py

@@ -0,0 +1,366 @@
+# SPDX-License-Identifier: Apache-2.0.
+# Copyright (c) 2024 - 2025 Waldiez and contributors.
+# pylint: disable=too-few-public-methods,no-self-use
+
+# before agent: query engine initialization
+# agent extra args:
+#   parsed_docs_path: Optional[Union[str, Path]] = None, # (we use str)
+#   collection_name: Optional[str] = None,  # (we use str)
+#   query_engine: Optional[RAGQueryEngine] = None,  # (we use RAGQueryEngine)
+
+"""Document Agent Extras."""
+
+# import shutil
+import shutil
+from pathlib import Path
+from typing import Optional
+
+from waldiez.exporting.core.extras.path_resolver import DefaultPathResolver
+from waldiez.exporting.core.protocols import PathResolver
+from waldiez.exporting.core.utils.llm_config import get_agent_llm_config_arg
+from waldiez.models import (
+    WaldiezAgent,
+    WaldiezDocAgent,
+    WaldiezDocAgentQueryEngine,
+    WaldiezModel,
+)
+
+from ...core import (
+    CodeExecutionConfig,
+    DefaultSerializer,
+    ImportPosition,
+    ImportStatement,
+    InstanceArgument,
+    Serializer,
+    SystemMessageConfig,
+    TerminationConfig,
+)
+from ...core.extras.agent_extras import StandardExtras
+
+
+class DocAgentProcessor:
+    """Process document agent requests."""
+
+    extras: StandardExtras
+
+    def __init__(
+        self,
+        agent: WaldiezAgent,
+        agent_names: dict[str, str],
+        model_names: dict[str, str],
+        all_models: list[WaldiezModel],
+        cache_seed: Optional[int] = None,
+        serializer: Serializer | None = None,
+        path_resolver: PathResolver | None = None,
+        output_dir: Optional[Path] = None,
+    ):
+        self.agent = agent
+        self.agent_name = agent_names.get(agent.id, agent.name)
+        self.output_dir = output_dir
+        self.model_names = model_names
+        self.all_models = all_models
+        self.cache_seed = cache_seed
+        self.path_resolver = path_resolver or DefaultPathResolver()
+        self.serializer = serializer or DefaultSerializer()
+
+    def process(
+        self,
+        code_execution_config: Optional[CodeExecutionConfig] = None,
+        termination_config: Optional[TerminationConfig] = None,
+        system_message_config: Optional[SystemMessageConfig] = None,
+    ) -> StandardExtras:
+        """Process the document agent configuration.
+
+        Parameters
+        ----------
+        code_execution_config : CodeExecutionConfig, optional
+            Configuration for code execution, if applicable.
+        termination_config : TerminationConfig, optional
+            Configuration for termination, if applicable.
+        system_message_config : SystemMessageConfig, optional
+            Configuration for system messages, if applicable.
+
+        Returns
+        -------
+        DocExtras
+            The document agent extras.
+        """
+        self.extras = StandardExtras(
+            instance_id=self.agent.id,
+            code_execution_config=code_execution_config,
+            termination_config=termination_config,
+            system_message_config=system_message_config,
+        )
+        if not self.agent.is_doc_agent or not isinstance(
+            self.agent, WaldiezDocAgent
+        ):
+            return self.extras
+        parsed_docs_path = self.get_parsed_docs_path()
+        self.extras.add_arg(
+            InstanceArgument(
+                instance_id=self.agent.id,
+                name="parsed_docs_path",
+                value=parsed_docs_path,
+                tabs=1,
+            )
+        )
+        self._get_query_engine_extras()
+        if self.extras.before_agent:
+            self.extras.add_arg(
+                InstanceArgument(
+                    instance_id=self.agent.id,
+                    name="query_engine",
+                    value=f"{self.agent_name}_query_engine",
+                    tabs=1,
+                )
+            )
+            self.extras.prepend_before_agent(self.extras.before_agent)
+        return self.extras
+
+    def _get_query_engine_extras(self) -> None:
+        if not isinstance(self.agent, WaldiezDocAgent):
+            return
+        query_engine = self.agent.get_query_engine()
+        if query_engine.type == "InMemoryQueryEngine":
+            self.get_in_memory_query_engine_extras()
+        elif query_engine.type == "VectorChromaQueryEngine":
+            self.get_vector_chroma_query_engine_extras(
+                self.agent.data.get_collection_name(),
+            )
+        elif (
+            query_engine.type == "VectorChromaCitationQueryEngine"
+        ):  # pragma: no branch
+            self.get_vector_chroma_citation_query_engine_extras(
+                query_engine,
+                self.agent.data.get_collection_name(),
+            )
+
+    def get_in_memory_query_engine_extras(self) -> None:
+        """Get the in-memory query engine extras."""
+        # InMemoryQueryEngine(llm_config=llm_config)
+        self.extras.add_import(
+            ImportStatement(
+                statement=(
+                    "from autogen.agents.experimental.document_agent."
+                    "inmemory_query_engine import InMemoryQueryEngine"
+                )
+            )
+        )
+        llm_config_arg = get_agent_llm_config_arg(
+            agent=self.agent,
+            tabs=1,
+            model_names=self.model_names,
+            all_models=self.all_models,
+            cache_seed=self.cache_seed,
+        )
+        self.extras.before_agent = (
+            f"{self.agent_name}_query_engine = "
+            f"InMemoryQueryEngine(\n{llm_config_arg}\n)"
+        )
+
+    def get_vector_chroma_query_engine_extras(
+        self,
+        collection_name: str,
+    ) -> None:
+        """Get the vector Chroma query engine extras.
+
+        Parameters
+        ----------
+        collection_name : str
+            The name of the collection to use.
+        """
+        # VectorChromaQueryEngine(
+        #   db_path: Optional[str] = None,
+        #   embedding_function: "Optional[EmbeddingFunction[Any]]" = None,
+        #   metadata: Optional[dict[str, Any]] = None,
+        #   llm: Optional["LLM"] = None,
+        #   collection_name: Optional[str] = None,
+        # )
+        self.extras.add_import(
+            ImportStatement(
+                statement=(
+                    "from autogen.agents.experimental.document_agent."
+                    "chroma_query_engine import VectorChromaQueryEngine"
+                )
+            )
+        )
+        llm_arg, before = self.get_llm_arg()
+        db_path = self.get_db_path()
+        self.extras.before_agent = before + "\n"
+        q_engine_init = (
+            f"{self.agent_name}_query_engine = "
+            f"VectorChromaQueryEngine(\n    {llm_arg}"
+        )
+        if db_path:  # pragma: no branch
+            q_engine_init += f",\n    db_path={db_path}"
+        if collection_name:  # pragma: no branch
+            q_engine_init += f',\n    collection_name="{collection_name}"'
+        q_engine_init += ",\n)"
+        self.extras.before_agent += q_engine_init
+
+    def get_vector_chroma_citation_query_engine_extras(
+        self,
+        query_engine: WaldiezDocAgentQueryEngine,
+        collection_name: str,
+    ) -> None:
+        """Get the vector Chroma citation query engine extras.
+
+        Parameters
+        ----------
+        query_engine : WaldiezDocAgentQueryEngine
+            The vector Chroma citation query engine to process.
+        collection_name : str
+            The name of the collection to use.
+        """
+        # VectorChromaCitationQueryEngine(
+        #   llm: Optional["LLM"] = None,
+        #   collection_name: Optional[str] = None,
+        #   enable_query_citations: bool = False,
+        #   citation_chunk_size: int = 512,
+        # )
+        self.extras.add_import(
+            ImportStatement(
+                statement=(
+                    "from autogen.agents.experimental.document_agent."
+                    "chroma_query_engine import VectorChromaCitationQueryEngine"
+                )
+            )
+        )
+        db_path = self.get_db_path()
+        llm_arg, before = self.get_llm_arg()
+        self.extras.before_agent = before + "\n"
+        q_engine_init = (
+            f"{self.agent_name}_query_engine = "
+            f"VectorChromaCitationQueryEngine({llm_arg}"
+        )
+        if db_path:  # pragma: no branch
+            q_engine_init += f",\n    db_path={db_path}"
+        if collection_name:  # pragma: no branch
+            q_engine_init += f',\n    collection_name="{collection_name}"'
+        enable_query_citations = query_engine.enable_query_citations
+        chunk_size = query_engine.citation_chunk_size or 512
+        q_engine_init += (
+            f",\n    enable_query_citations={enable_query_citations}"
+            f",\n    citation_chunk_size={chunk_size}"
+        )
+        q_engine_init += ",\n)"
+        self.extras.before_agent += q_engine_init
+
+    def get_llm_arg(self) -> tuple[str, str]:
+        """Get the LLM argument for the agent and any content before it.
+
+        Returns
+        -------
+        str
+            The LLM argument for the agent and any content before it.
+        """
+        agent_model = get_agent_model(
+            agent=self.agent,
+            all_models=self.all_models,
+        )
+        if not agent_model:
+            # ag2's default
+            self.extras.add_import(
+                ImportStatement(
+                    statement="from llama_index.llms.openai import OpenAI",
+                )
+            )
+            return 'llm=OpenAI(model="gpt-4o", temperature=0.0)', ""
+        arg, before = agent_model.get_llm_arg()
+        for import_stmt in agent_model.get_llm_imports():
+            self.extras.add_import(
+                ImportStatement(
+                    statement=import_stmt,
+                    position=ImportPosition.THIRD_PARTY,
+                )
+            )
+        return f"llm={arg}", before
+
+    def get_db_path(self) -> str:
+        """Get the database path for the document agent.
+
+        Returns
+        -------
+        str
+            The database path for the document agent.
+
+        Raises
+        ------
+        TypeError
+            If the agent is not a WaldiezDocAgent instance.
+        """
+        if not isinstance(self.agent, WaldiezDocAgent):  # pragma: no cover
+            raise TypeError("Agent must be a WaldiezDocAgent instance.")
+        query_engine = self.agent.get_query_engine()
+        db_path_str = ""
+        if query_engine.get_db_path():
+            db_path_str = query_engine.get_db_path()
+        else:
+            if self.output_dir:
+                chroma_dir = self.output_dir / "chroma"
+                chroma_dir.mkdir(parents=True, exist_ok=True)
+                db_path_str = str(chroma_dir)
+            else:
+                db_path_str = query_engine.get_db_path()
+        if self.agent.data.reset_collection is True:
+            shutil.rmtree(db_path_str, ignore_errors=True)
+            Path(db_path_str).mkdir(parents=True, exist_ok=True)
+        return self.path_resolver.resolve(db_path_str)
+
+    def get_parsed_docs_path(self) -> str:
+        """Get the parsed documents path for the document agent.
+
+        Returns
+        -------
+        str
+            The parsed documents path for the document agent.
+
+        Raises
+        ------
+        TypeError
+            If the agent is not a WaldiezDocAgent instance.
+        """
+        if not isinstance(self.agent, WaldiezDocAgent):
+            raise TypeError("Agent must be a WaldiezDocAgent instance.")
+        parsed_docs_path = self.agent.data.parsed_docs_path or ""
+        if not parsed_docs_path:
+            if self.output_dir:
+                parsed_docs_path = str(self.output_dir / "parsed_docs")
+            else:
+                parsed_docs_path = self.agent.get_parsed_docs_path()
+        resolved = Path(parsed_docs_path).resolve()
+        if not resolved.is_absolute():
+            parsed_docs_path = str(Path.cwd() / parsed_docs_path)
+        if not Path(parsed_docs_path).is_dir():
+            Path(parsed_docs_path).mkdir(parents=True, exist_ok=True)
+        if self.agent.data.reset_collection is True:
+            shutil.rmtree(parsed_docs_path, ignore_errors=True)
+            Path(parsed_docs_path).mkdir(parents=True, exist_ok=True)
+        return self.path_resolver.resolve(str(parsed_docs_path))
+
+
+def get_agent_model(
+    agent: WaldiezAgent,
+    all_models: list[WaldiezModel],
+) -> WaldiezModel | None:
+    """Get the first model of the agent.
+
+    Parameters
+    ----------
+    agent : WaldiezAgent
+        The agent.
+    all_models : list[WaldiezModel]
+        All the models in the flow.
+
+    Returns
+    -------
+    WaldiezModel | None
+        The first model of the agent, or None if not found.
+    """
+    if not agent.data.model_ids:
+        return None
+    for model_id in agent.data.model_ids:
+        model = next((m for m in all_models if m.id == model_id), None)
+        if model is not None:
+            return model
+    return None

waldiez/exporting/agent/extras/group_member_extras.py

@@ -68,8 +68,9 @@ class GroupMemberAgentProcessor:
         self.result = GroupMemberProcessorResult()
         if not self.agent.is_group_member:
             return self.result
-        self._process_agent_functions()
-        self._process_agent_update_state_before_reply()
+        if not self.agent.is_doc_agent:
+            self._process_agent_functions()
+            self._process_agent_update_state_before_reply()
         self._process_handoff_registrations()
         return self.result