ag2 0.10.2__py3-none-any.whl

autogen/agents/experimental/document_agent/document_agent.py

@@ -0,0 +1,643 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from copy import deepcopy
+from pathlib import Path
+from typing import Annotated, Any, cast
+
+from pydantic import BaseModel, Field
+
+from .... import Agent, ConversableAgent, UpdateSystemMessage
+from ....agentchat.contrib.rag.query_engine import RAGQueryEngine
+from ....agentchat.group.context_condition import ExpressionContextCondition
+from ....agentchat.group.context_expression import ContextExpression
+from ....agentchat.group.context_variables import ContextVariables
+from ....agentchat.group.llm_condition import StringLLMCondition
+from ....agentchat.group.multi_agent_chat import initiate_group_chat
+from ....agentchat.group.on_condition import OnCondition
+from ....agentchat.group.on_context_condition import OnContextCondition
+from ....agentchat.group.patterns.pattern import DefaultPattern
+from ....agentchat.group.reply_result import ReplyResult
+from ....agentchat.group.targets.transition_target import AgentNameTarget, AgentTarget, StayTarget, TerminateTarget
+from ....doc_utils import export_module
+from ....llm_config import LLMConfig
+from ....oai.client import OpenAIWrapper
+from .chroma_query_engine import VectorChromaQueryEngine
+from .docling_doc_ingest_agent import DoclingDocIngestAgent
+from .document_conditions import SummaryTaskAvailableCondition
+from .document_utils import Ingest, Query
+
+__all__ = ["DocAgent"]
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SYSTEM_MESSAGE = """
+    You are a document agent.
+    You are given a list of documents to ingest and a list of queries to perform.
+    You are responsible for ingesting the documents and answering the queries.
+"""
+TASK_MANAGER_NAME = "TaskManagerAgent"
+TASK_MANAGER_SYSTEM_MESSAGE = """
+    You are a task manager agent. You have 2 priorities:
+    1. You initiate the tasks which updates the context variables based on the task decisions (DocumentTask) from the DocumentTriageAgent.
+    ALWAYS call initiate_tasks first when you receive a message from the DocumentTriageAgent, even if you think there are no new tasks.
+    This ensures that any new ingestions or queries from the triage agent are properly recorded.
+    Put all ingestion and query tasks into the one tool call.
+        i.e. output
+        {
+            "ingestions": [
+                {
+                    "path_or_url": "path_or_url"
+                }
+            ],
+            "queries": [
+                {
+                    "query_type": "RAG_QUERY",
+                    "query": "query"
+                }
+            ],
+            "query_results": [
+                {
+                    "query": "query",
+                    "result": "result"
+                }
+            ]
+        }
+    2. If there are no documents to ingest and no queries to run, hand control off to the summary agent.
+
+    Put all file paths and URLs into the ingestions. A http/https URL is also a valid path and should be ingested.
+
+    Use the initiate_tasks tool to incorporate all ingestions and queries. Don't call it again until new ingestions or queries are raised.
+
+    New ingestions and queries may be raised from time to time, so use the initiate_tasks again if you see new ingestions/queries.
+
+    Once the ingestion and query tasks are done, return to the summary agent, do not ask questions of the user.
+    """
+
+DEFAULT_ERROR_GROUP_CHAT_MESSAGE: str = """
+Document Agent failed to perform task.
+"""
+
+ERROR_MANAGER_NAME = "ErrorManagerAgent"
+ERROR_MANAGER_SYSTEM_MESSAGE = """
+You communicate errors to the user. Include the original error messages in full. Use the format:
+The following error(s) have occurred:
+- Error 1
+- Error 2
+"""
+
+
+class DocumentTask(BaseModel):
+    """The structured output format for task decisions."""
+
+    ingestions: list[Ingest] = Field(description="The list of documents to ingest.")
+    queries: list[Query] = Field(description="The list of queries to perform.")
+
+    def format(self) -> str:
+        """Format the DocumentTask as a string for the TaskManager to work with."""
+        if len(self.ingestions) == 0 and len(self.queries) == 0:
+            return "There were no ingestion or query tasks detected."
+
+        instructions = "Tasks:\n\n"
+        order = 1
+
+        if len(self.ingestions) > 0:
+            instructions += "Ingestions:\n"
+            for ingestion in self.ingestions:
+                instructions += f"{order}: {ingestion.path_or_url}\n"
+                order += 1
+
+            instructions += "\n"
+
+        if len(self.queries) > 0:
+            instructions += "Queries:\n"
+            for query in self.queries:
+                instructions += f"{order}: {query.query}\n"
+                order += 1
+
+        return instructions
+
+
+class DocumentTriageAgent(ConversableAgent):
+    """The DocumentTriageAgent is responsible for deciding what type of task to perform from user requests."""
+
+    def __init__(self, llm_config: LLMConfig | dict[str, Any] | None = None):
+        # Add the structured message to the LLM configuration
+        structured_config_list = deepcopy(llm_config)
+        structured_config_list["response_format"] = DocumentTask  # type: ignore[index]
+
+        super().__init__(
+            name="DocumentTriageAgent",
+            system_message=(
+                "You are a document triage agent. "
+                "You are responsible for deciding what type of task to perform from a user's request and populating a DocumentTask formatted response. "
+                "If the user specifies files or URLs, add them as individual 'ingestions' to DocumentTask. "
+                "You can access external websites if given a URL, so put them in as ingestions. "
+                "Add the user's questions about the files/URLs as individual 'RAG_QUERY' queries to the 'query' list in the DocumentTask. "
+                "Don't make up questions, keep it as concise and close to the user's request as possible."
+            ),
+            human_input_mode="NEVER",
+            llm_config=structured_config_list,
+        )
+
+
+@export_module("autogen.agents.experimental")
+class DocAgent(ConversableAgent):
+    """The DocAgent is responsible for ingest and querying documents.
+
+    Internally, it generates a group chat with a set of agents to ingest, query, and summarize.
+    """
+
+    def __init__(
+        self,
+        name: str | None = None,
+        llm_config: LLMConfig | dict[str, Any] | None = None,
+        system_message: str | None = None,
+        parsed_docs_path: str | Path | None = None,
+        collection_name: str | None = None,
+        query_engine: RAGQueryEngine | None = None,
+    ):
+        """Initialize the DocAgent.
+
+        Args:
+            name (Optional[str]): The name of the DocAgent.
+            llm_config (Optional[LLMConfig, dict[str, Any]]): The configuration for the LLM.
+            system_message (Optional[str]): The system message for the DocAgent.
+            parsed_docs_path (Union[str, Path]): The path where parsed documents will be stored.
+            collection_name (Optional[str]): The unique name for the data store collection. If omitted, a random name will be used. Populate this to reuse previous ingested data.
+            query_engine (Optional[RAGQueryEngine]): The query engine to use for querying documents, defaults to VectorChromaQueryEngine if none provided.
+                                                     Use enable_query_citations and implement query_with_citations method to enable citation support. e.g. VectorChromaCitationQueryEngine
+
+        The DocAgent is responsible for generating a group of agents to solve a task.
+
+        The agents that the DocAgent generates are:
+        - Triage Agent: responsible for deciding what type of task to perform from user requests.
+        - Task Manager Agent: responsible for managing the tasks.
+        - Parser Agent: responsible for parsing the documents.
+        - Data Ingestion Agent: responsible for ingesting the documents.
+        - Query Agent: responsible for answering the user's questions.
+        - Error Agent: responsible for returning errors gracefully.
+        - Summary Agent: responsible for generating a summary of the user's questions.
+        """
+        name = name or "DocAgent"
+        llm_config = llm_config or LLMConfig.get_current_llm_config()
+        system_message = system_message or DEFAULT_SYSTEM_MESSAGE
+        parsed_docs_path = parsed_docs_path or "./parsed_docs"
+
+        # Default Query Engine will be ChromaDB
+        if query_engine is None:
+            query_engine = VectorChromaQueryEngine(collection_name=collection_name)
+
+        super().__init__(
+            name=name,
+            system_message=system_message,
+            llm_config=llm_config,
+            human_input_mode="NEVER",
+        )
+        self.register_reply([ConversableAgent, None], self.generate_inner_group_chat_reply, position=0)
+
+        self._triage_agent = DocumentTriageAgent(llm_config=llm_config)
+
+        def create_error_agent_prompt(agent: ConversableAgent, messages: list[dict[str, Any]]) -> str:
+            """Create the error agent prompt, primarily used to update ingested documents for ending.
+
+            Args:
+                agent: The conversable agent requesting the prompt
+                messages: List of conversation messages
+
+            Returns:
+                str: The error manager system message
+            """
+            update_ingested_documents()
+
+            return ERROR_MANAGER_SYSTEM_MESSAGE
+
+        self._error_agent = ConversableAgent(
+            name=ERROR_MANAGER_NAME,
+            system_message=ERROR_MANAGER_SYSTEM_MESSAGE,
+            llm_config=llm_config,
+            update_agent_state_before_reply=[UpdateSystemMessage(create_error_agent_prompt)],
+        )
+
+        def update_ingested_documents() -> None:
+            """Updates the list of ingested documents, persisted so we can keep a list over multiple replies.
+
+            This function updates self.documents_ingested with any new documents that have been ingested
+            by the triage agent, ensuring persistence across multiple DocAgent interactions.
+            """
+            agent_documents_ingested = self._triage_agent.context_variables.get("DocumentsIngested", [])
+            # Update self.documents_ingested with any new documents ingested
+            for doc in agent_documents_ingested:  # type: ignore[union-attr]
+                if doc not in self.documents_ingested:
+                    self.documents_ingested.append(doc)
+
+        class TaskInitInfo(BaseModel):
+            ingestions: Annotated[list[Ingest], Field(description="List of documents, files, and URLs to ingest")]
+            queries: Annotated[list[Query], Field(description="List of queries to run")]
+
+        def _deduplicate_ingestions(
+            new_ingestions: list[Ingest], existing_ingestions: list[Ingest], documents_ingested: list[str]
+        ) -> tuple[list[Ingest], list[str]]:
+            """Deduplicate ingestions against existing pending and already ingested documents.
+
+            Args:
+                new_ingestions: List of new ingestion requests to process
+                existing_ingestions: List of ingestions already pending
+                documents_ingested: List of document paths already ingested
+
+            Returns:
+                tuple: (new_unique_ingestions, ignored_duplicate_paths)
+            """
+            unique_ingestions = []
+            ignored_paths = []
+
+            for ingestion in new_ingestions:
+                ingestion_path = ingestion.path_or_url
+                # Check if already in pending ingestions
+                already_pending = any(existing.path_or_url == ingestion_path for existing in existing_ingestions)
+                # Check if already ingested
+                already_ingested = ingestion_path in documents_ingested
+
+                if already_pending or already_ingested:
+                    ignored_paths.append(ingestion_path)
+                else:
+                    unique_ingestions.append(ingestion)
+
+            return unique_ingestions, ignored_paths
+
+        def _deduplicate_queries(
+            new_queries: list[Query], existing_queries: list[Query]
+        ) -> tuple[list[Query], list[str]]:
+            """Deduplicate queries against existing pending queries.
+
+            Args:
+                new_queries: List of new query requests to process
+                existing_queries: List of queries already pending
+
+            Returns:
+                tuple: (new_unique_queries, ignored_duplicate_query_texts)
+            """
+            unique_queries = []
+            ignored_query_texts = []
+
+            for query in new_queries:
+                query_text = query.query
+                # Check if query already exists in pending queries
+                already_pending = any(existing.query == query_text for existing in existing_queries)
+
+                if already_pending:
+                    ignored_query_texts.append(query_text)
+                else:
+                    unique_queries.append(query)
+
+            return unique_queries, ignored_query_texts
+
+        def _build_response_message(
+            added_ingestions: int, ignored_ingestions: list[str], added_queries: int, ignored_queries: list[str]
+        ) -> str:
+            """Build a descriptive response message about what was added/ignored.
+
+            Args:
+                added_ingestions: Number of unique ingestions added
+                ignored_ingestions: List of duplicate ingestion paths ignored
+                added_queries: Number of unique queries added
+                ignored_queries: List of duplicate query texts ignored
+
+            Returns:
+                str: Formatted message describing the results
+            """
+            messages = []
+
+            if added_ingestions > 0:
+                messages.append(f"Added {added_ingestions} new document(s) for ingestion")
+
+            if ignored_ingestions:
+                messages.append(
+                    f"Ignored {len(ignored_ingestions)} duplicate document(s): {', '.join(ignored_ingestions)}"
+                )
+
+            if added_queries > 0:
+                messages.append(f"Added {added_queries} new query/queries")
+
+            if ignored_queries:
+                messages.append(f"Ignored {len(ignored_queries)} duplicate query/queries: {', '.join(ignored_queries)}")
+
+            if messages:
+                return "; ".join(messages)
+            else:
+                return "All requested tasks were duplicates and ignored"
+
+        def initiate_tasks(
+            task_init_info: Annotated[TaskInitInfo, "Documents, Files, URLs to ingest and the queries to run"],
+            context_variables: Annotated[ContextVariables, "Context variables"],
+        ) -> ReplyResult:
+            """Add documents to ingest and queries to answer when received.
+
+            Args:
+                task_init_info: Information about documents to ingest and queries to run
+                context_variables: The current context variables containing task state
+
+            Returns:
+                ReplyResult: Contains response message, updated context, and target agent
+            """
+            ingestions = task_init_info.ingestions
+            queries = task_init_info.queries
+
+            if "TaskInitiated" in context_variables:
+                # Handle follow-up tasks with deduplication
+                added_ingestions_count = 0
+                ignored_ingestions = []
+                added_queries_count = 0
+                ignored_queries = []
+
+                if ingestions:
+                    existing_ingestions: list[Ingest] = context_variables.get("DocumentsToIngest", [])  # type: ignore[assignment]
+                    documents_ingested: list[str] = context_variables.get("DocumentsIngested", [])  # type: ignore[assignment]
+
+                    unique_ingestions, ignored_ingestion_paths = _deduplicate_ingestions(
+                        ingestions, existing_ingestions, documents_ingested
+                    )
+
+                    if unique_ingestions:
+                        context_variables["DocumentsToIngest"] = existing_ingestions + unique_ingestions
+                        added_ingestions_count = len(unique_ingestions)
+
+                    ignored_ingestions = ignored_ingestion_paths
+
+                if queries:
+                    existing_queries: list[Query] = context_variables.get("QueriesToRun", [])  # type: ignore[assignment]
+
+                    unique_queries, ignored_query_texts = _deduplicate_queries(queries, existing_queries)
+
+                    if unique_queries:
+                        context_variables["QueriesToRun"] = existing_queries + unique_queries
+                        added_queries_count = len(unique_queries)
+
+                    ignored_queries = ignored_query_texts
+
+                if not ingestions and not queries:
+                    return ReplyResult(message="No new tasks to initiate", context_variables=context_variables)
+
+                response_message = _build_response_message(
+                    added_ingestions_count, ignored_ingestions, added_queries_count, ignored_queries
+                )
+
+            else:
+                # First time initialization - no deduplication needed
+                context_variables["DocumentsToIngest"] = ingestions
+                context_variables["QueriesToRun"] = list(queries)
+                context_variables["TaskInitiated"] = True
+                response_message = "Updated context variables with task decisions"
+
+            return ReplyResult(
+                message=response_message,
+                context_variables=context_variables,
+                target=AgentNameTarget(agent_name=TASK_MANAGER_NAME),
+            )
+
+        self._task_manager_agent = ConversableAgent(
+            name=TASK_MANAGER_NAME,
+            system_message=TASK_MANAGER_SYSTEM_MESSAGE,
+            llm_config=llm_config,
+            functions=[initiate_tasks],
+        )
+
+        self._triage_agent.handoffs.set_after_work(target=AgentTarget(agent=self._task_manager_agent))
+
+        self._data_ingestion_agent = DoclingDocIngestAgent(
+            llm_config=llm_config,
+            query_engine=query_engine,
+            parsed_docs_path=parsed_docs_path,
+            return_agent_success=TASK_MANAGER_NAME,
+            return_agent_error=ERROR_MANAGER_NAME,
+        )
+
+        def execute_rag_query(context_variables: ContextVariables) -> ReplyResult:  # type: ignore[type-arg]
+            """Execute outstanding RAG queries, call the tool once for each outstanding query. Call this tool with no arguments.
+
+            Args:
+                context_variables: The current context variables containing queries to run
+
+            Returns:
+                ReplyResult: Contains query answer, updated context, and target agent
+            """
+            if len(context_variables["QueriesToRun"]) == 0:
+                return ReplyResult(
+                    target=AgentNameTarget(agent_name=TASK_MANAGER_NAME),
+                    message="No queries to run",
+                    context_variables=context_variables,
+                )
+
+            query = context_variables["QueriesToRun"][0].query
+            try:
+                if (
+                    hasattr(query_engine, "enable_query_citations")
+                    and query_engine.enable_query_citations
+                    and hasattr(query_engine, "query_with_citations")
+                    and callable(query_engine.query_with_citations)
+                ):
+                    answer_with_citations = query_engine.query_with_citations(query)  # type: ignore[union-attr]
+                    answer = answer_with_citations.answer
+                    txt_citations = [
+                        {
+                            "text_chunk": source.node.get_text(),
+                            "file_path": source.metadata["file_path"],
+                        }
+                        for source in answer_with_citations.citations
+                    ]
+                    logger.info(f"Citations:\n {txt_citations}")
+                else:
+                    answer = query_engine.query(query)
+                    txt_citations = []
+                context_variables["QueriesToRun"].pop(0)
+                context_variables["CompletedTaskCount"] += 1
+                context_variables["QueryResults"].append({"query": query, "answer": answer, "citations": txt_citations})
+
+                # Query completed
+
+                return ReplyResult(message=answer, context_variables=context_variables)
+            except Exception as e:
+                return ReplyResult(
+                    target=AgentNameTarget(agent_name=ERROR_MANAGER_NAME),
+                    message=f"Query failed for '{query}': {e}",
+                    context_variables=context_variables,
+                )
+
+        self._query_agent = ConversableAgent(
+            name="QueryAgent",
+            system_message="""
+            You are a query agent.
+            You answer the user's questions only using the query function provided to you.
+            You can only call use the execute_rag_query tool once per turn.
+            """,
+            llm_config=llm_config,
+            functions=[execute_rag_query],
+        )
+
+        # Summary agent prompt will include the results of the ingestions and queries
+        def create_summary_agent_prompt(agent: ConversableAgent, messages: list[dict[str, Any]]) -> str:
+            """Create the summary agent prompt and updates ingested documents.
+
+            Args:
+                agent: The conversable agent requesting the prompt
+                messages: List of conversation messages
+
+            Returns:
+                str: The summary agent system message with context information
+            """
+            update_ingested_documents()
+
+            documents_to_ingest: list[Ingest] = cast(list[Ingest], agent.context_variables.get("DocumentsToIngest", []))
+            queries_to_run: list[Query] = cast(list[Query], agent.context_variables.get("QueriesToRun", []))
+
+            system_message = (
+                "You are a summary agent and you provide a summary of all completed tasks and the list of queries and their answers. "
+                "Output two sections: 'Ingestions:' and 'Queries:' with the results of the tasks. Number the ingestions and queries. "
+                "If there are no ingestions output 'No ingestions', if there are no queries output 'No queries' under their respective sections. "
+                "Don't add markdown formatting. "
+                "For each query, there is one answer and, optionally, a list of citations."
+                "For each citation, it contains two fields: 'text_chunk' and 'file_path'."
+                "Format the Query and Answers and Citations as 'Query:\nAnswer:\n\nCitations:'. Add a number to each query if more than one. Use the context below:\n"
+                "For each query, output the full citation contents and list them one by one,"
+                "format each citation as '\nSource [X] (chunk file_path here):\n\nChunk X:\n(text_chunk here)' and mark a separator between each citation using '\n#########################\n\n'."
+                "If there are no citations at all, DON'T INCLUDE ANY mention of citations.\n"
+                f"Documents ingested: {documents_to_ingest}\n"
+                f"Documents left to ingest: {len(documents_to_ingest)}\n"
+                f"Queries left to run: {len(queries_to_run)}\n"
+                f"Query and Answers and Citations: {queries_to_run}\n"
+            )
+
+            return system_message
+
+        self._summary_agent = ConversableAgent(
+            name="SummaryAgent",
+            llm_config=llm_config,
+            update_agent_state_before_reply=[UpdateSystemMessage(create_summary_agent_prompt)],
+        )
+
+        self._task_manager_agent.register_handoffs([
+            OnContextCondition(  # Go straight to data ingestion agent if we have documents to ingest
+                target=AgentTarget(agent=self._data_ingestion_agent),
+                condition=ExpressionContextCondition(
+                    expression=ContextExpression(expression="len(${DocumentsToIngest}) > 0")
+                ),
+            ),
+            OnContextCondition(  # Go to Query agent if we have queries to run (ingestion above run first)
+                target=AgentTarget(agent=self._query_agent),
+                condition=ExpressionContextCondition(
+                    expression=ContextExpression(expression="len(${QueriesToRun}) > 0")
+                ),
+            ),
+            # Removed automatic context condition - let task manager decide when to summarize
+            OnCondition(
+                target=AgentTarget(agent=self._summary_agent),
+                condition=StringLLMCondition(
+                    prompt="Call this function if all work is done and a summary will be created"
+                ),
+                available=SummaryTaskAvailableCondition(),  # Custom AvailableCondition class
+            ),
+        ])
+        self._task_manager_agent.handoffs.set_after_work(target=StayTarget())
+
+        self._data_ingestion_agent.handoffs.set_after_work(target=AgentTarget(agent=self._task_manager_agent))
+
+        self._query_agent.handoffs.set_after_work(target=AgentTarget(agent=self._task_manager_agent))
+
+        # Summary agent terminates the DocumentAgent
+        self._summary_agent.handoffs.set_after_work(target=TerminateTarget())
+
+        # The Error Agent always terminates the DocumentAgent
+        self._error_agent.handoffs.set_after_work(target=TerminateTarget())
+
+        self.register_reply([Agent, None], DocAgent.generate_inner_group_chat_reply)
+
+        self.documents_ingested: list[str] = []
+        self._group_chat_context_variables: ContextVariables | None = None
+
+    def generate_inner_group_chat_reply(
+        self,
+        messages: list[dict[str, Any]] | str | None = None,
+        sender: Agent | None = None,
+        config: OpenAIWrapper | None = None,
+    ) -> tuple[bool, str | dict[str, Any] | None]:
+        """Reply function that generates the inner group chat reply for the DocAgent.
+
+        Args:
+            messages: Input messages to process
+            sender: The agent that sent the message
+            config: OpenAI wrapper configuration
+
+        Returns:
+            tuple: (should_terminate, reply_message)
+        """
+        # Use existing context_variables if available, otherwise create new ones
+        if hasattr(self, "_group_chat_context_variables") and self._group_chat_context_variables is not None:
+            context_variables = self._group_chat_context_variables
+            # Reset for the new run
+            context_variables["DocumentsToIngest"] = []  # type: ignore[index]
+        else:
+            context_variables = ContextVariables(
+                data={
+                    "CompletedTaskCount": 0,
+                    "DocumentsToIngest": [],
+                    "DocumentsIngested": self.documents_ingested,
+                    "QueriesToRun": [],
+                    "QueryResults": [],
+                }
+            )
+            self._group_chat_context_variables = context_variables
+
+        group_chat_agents = [
+            self._triage_agent,
+            self._task_manager_agent,
+            self._data_ingestion_agent,
+            self._query_agent,
+            self._summary_agent,
+            self._error_agent,
+        ]
+
+        agent_pattern = DefaultPattern(
+            initial_agent=self._triage_agent,
+            agents=group_chat_agents,
+            context_variables=context_variables,
+            group_after_work=TerminateTarget(),
+        )
+
+        chat_result, context_variables, last_speaker = initiate_group_chat(
+            pattern=agent_pattern,
+            messages=self._get_document_input_message(messages),
+        )
+        if last_speaker == self._error_agent:
+            # If we finish with the error agent, we return their message which contains the error
+            return True, chat_result.summary
+        if last_speaker != self._summary_agent:
+            # If the group chat finished but not with the summary agent, we assume something has gone wrong with the flow
+            return True, DEFAULT_ERROR_GROUP_CHAT_MESSAGE
+
+        return True, chat_result.summary
+
+    def _get_document_input_message(self, messages: list[dict[str, Any]] | str | None) -> str:  # type: ignore[type-arg]
+        """Gets and validates the input message(s) for the document agent.
+
+        Args:
+            messages: Input messages as string or list of message dictionaries
+
+        Returns:
+            str: The extracted message content
+
+        Raises:
+            NotImplementedError: If messages format is invalid
+        """
+        if isinstance(messages, str):
+            return messages
+        elif (
+            isinstance(messages, list)
+            and len(messages) > 0
+            and "content" in messages[-1]
+            and isinstance(messages[-1]["content"], str)
+        ):
+            return messages[-1]["content"]
+        else:
+            raise NotImplementedError("Invalid messages format. Must be a list of messages or a string.")