Agents

def process(self) -> CodeExecutionConfig:
    """Process code execution configuration.

    Returns
    -------
    CodeExecutionConfig
        The processed code execution configuration.
    """
    if self.agent.data.code_execution_config is False:
        return CodeExecutionConfig(
            executor_content="",
            executor_argument="False",
            executor_import=None,
        )

    config = self.agent.data.code_execution_config
    use_docker = (
        config.use_docker if config.use_docker is not None else False
    )

    executor_class = self._get_executor_class_name(use_docker)
    executor_content = self._build_executor_content(config, use_docker)
    executor_arg = f'{{"executor": {self.agent_name}_executor}}'
    executor_import = f"from autogen.coding import {executor_class}"

    return CodeExecutionConfig(
        executor_content=executor_content,
        executor_argument=executor_arg,
        executor_import=ImportStatement(
            statement=executor_import,
            position=ImportPosition.THIRD_PARTY,
        ),
    )

def create_reasoning_extras(self) -> StandardExtras:
    """Create reasoning-specific extras for the agent.

    Returns
    -------
    StandardExtras
        The reasoning-specific extras.
    """
    # Start with standard extras
    extras = self._create_standard_extras()

    reasoning_processor = ReasoningAgentProcessor(
        agent=self.agent,
        serializer=(
            self.context.get_serializer()
            if self.context.serializer
            else None
        ),
    )
    # Process reasoning-specific content
    reason_extras = reasoning_processor.process(
        code_execution_config=extras.code_execution_config,
        termination_config=extras.termination_config,
        system_message_config=extras.system_message_config,
    )
    # Add reasoning arguments
    for arg in reason_extras.extra_args:
        extras.add_arg(arg)

    return reason_extras

def generate_main_content(self) -> str | None:
    """Generate the main agent definition.

    Returns
    -------
    Optional[str]
        The main agent definition content, or None if not applicable.
    """
    # Use AgentProcessor to generate the main definition
    agent_processor = AgentProcessor(
        agent=self.agent,
        agent_names=self.agent_names,
        arguments_resolver=self.arguments_resolver,
        extras=self.extras,
    )
    result = agent_processor.process().content
    if result:
        # Add the main agent definition content
        self.add_content(
            result,
            ExportPosition.AGENTS,
            order=(
                ContentOrder.MAIN_CONTENT
                if not self.agent.is_group_manager
                else ContentOrder.LATE_CLEANUP
            ),
        )
        self.extras.append_after_all_agents(result)
    return result

def process(
    self,
    code_execution_config: CodeExecutionConfig | None = None,
    termination_config: TerminationConfig | None = None,
    system_message_config: SystemMessageConfig | None = None,
) -> CaptainExtras:
    """Process group manager and return extras.

    Parameters
    ----------
    code_execution_config : CodeExecutionConfig, optional
        Configuration for code execution, by default None.
    termination_config : TerminationConfig, optional
        Configuration for termination, by default None.
    system_message_config : SystemMessageConfig, optional
        Configuration for system messages, by default None.

    Returns
    -------
    CaptainExtras
        The processed result containing extra arguments, before content,
        imports, and environment variables.
    """
    result = CaptainExtras(
        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_captain or not isinstance(
        self.agent, WaldiezCaptainAgent
    ):
        return result
    agent_name = self.agent_names[self.agent.id]
    save_path = str(self.output_dir) if self.output_dir else "."
    if save_path != ".":
        os.makedirs(save_path, exist_ok=True)
    result.add_arg("agent_config_save_path=os.getcwd()", tabs=1)
    if self.agent.data.agent_lib:
        lib_dict = [
            lib.model_dump(by_alias=False)
            for lib in self.agent.data.agent_lib
        ]
        lib_json_name = f"{agent_name}_agent_lib.json"
        agent_lib_path = os.path.join(save_path, lib_json_name)
        with open(agent_lib_path, "w", encoding="utf-8", newline="\n") as f:
            json.dump(lib_dict, f, ensure_ascii=False, indent=4)
        result.add_arg(f'agent_lib="{lib_json_name}"', tabs=1)
    if self.agent.data.tool_lib:
        result.add_arg(f'tool_lib="{self.agent.data.tool_lib}"', tabs=1)
    nested_config = self._generate_nested_config(
        self.agent,
        agent_name,
        save_path,
    )
    result.set_nested_config(nested_config)
    serialized_nested_config = self.serializer.serialize(nested_config)
    result.add_arg(f"nested_config={serialized_nested_config}", tabs=1)
    result.append_before_agent(
        self._patch_default_nested_config(),
    )
    return result

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()
    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:
        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_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_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_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:
        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_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_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