Repository: maihemlabs/maihem (language: python)

Issue:
simulator



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 217bf50..9b9ffd2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,4 +6,5 @@ venv/
 .vscode/
 myenv/
 **/myenv/
-build/
\ No newline at end of file
+build/
+src/maihem.egg-info/
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index e860ff3..b6b07ab 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,7 @@ classifiers = [
     "Programming Language :: Python :: 3",
 ]
 keywords = ["genAI", "LLM", "evaluation", "data", "generation", "synthetic"]
-dependencies = ["requests", "pandas", "pydantic", "pycountry", "pydantic_extra_types", "httpx", "attrs", "loguru", "tqdm", "anyio", "pymupdf", "docx2txt", "yaspin"]
+dependencies = ["requests", "pandas", "pydantic", "pycountry", "pydantic_extra_types", "httpx", "attrs", "loguru", "tqdm", "anyio", "pymupdf", "docx2txt", "yaspin", "pyyaml"]
 requires-python = ">=3.9"
 
 [project.urls]
diff --git a/src/config.yaml b/src/config.yaml
new file mode 100644
index 0000000..afeb3db
--- /dev/null
+++ b/src/config.yaml
@@ -0,0 +1,49 @@
+# See list of supported metrics here: https://docs.maihem.ai/metrics/collection
+metric: qa_cx_goal_completion
+
+# Agent that starts the conversation (maihem or target_agent)
+initiating_agent: maihem
+
+target_agent:
+  # Unique identifier for the target agent.
+  tg_ag_1_upsert_1:
+    # (Optional) Name of the target agent to be displayed in the Maihem webapp
+    name: Therapy chatbot 
+
+    # Free text short description of the role of the target agent
+    role: mental health assistant
+
+    # Free text description of the industry that the target agent belongs to
+    industry: mental health
+
+    # Free text detailed description of the target agent
+    description: A chatbot for emotional support
+
+    # Language of the target agent in ISO 639-1 format
+    language: en  
+
+  # Add more target agents if needed with a different unique identifier
+  tg_ag_2:
+    name: Therapy chatbot 2
+    role: mental health assistant
+    industry: mental health
+    description: A chatbot for emotional support and venting
+    language: en  
+
+maihem_agent:
+  # Unique identifier for the Maihem agent
+  mh_ag_1_upsert_1:
+    # (Optional) Name of the Maihem agent to be displayed in the Maihem webapp
+    name: Maihem_Agent_1
+
+    # Free text prompt to guide the behavior of the Maihem Agent
+    behavior_prompt: Vent about stress in the workplace
+
+    # Maximum number of turns for the conversation
+    conversation_turns_max: 4
+  
+  # Add more Maihem agents if needed with a different unique identifier
+  mh_ag_2:
+    name: Maihem_Agent_2
+    behavior_prompt: Discuss techniques for managing anxiety
+    conversation_turns_max: 10
\ No newline at end of file
diff --git a/src/design_maihem_new.py b/src/design_maihem_new.py
new file mode 100644
index 0000000..fcd9118
--- /dev/null
+++ b/src/design_maihem_new.py
@@ -0,0 +1,66 @@
+from maihem.agents import TargetAgent, MaihemAgent
+from maihem.clients import Maihem
+
+
+# Create Maihem Client
+maihem_client = Maihem()
+
+# Create Target Agent
+target_agent = maihem_client.create_target_agent(
+    name="PG_Chatbot_1", # Name of the Target Agent, does not have to be unique
+    identifier="tg_ag_1234567", # (Optional) Define unique identifier for your Target Agent, or we create one for you
+    role="mental health assistant", # Enum
+    industry="mental health", # Enum
+    description="A chatbot for emotional support", # Free text
+    language="en", # (Optional) Language of the Target Agent, default is "en"
+)
+
+target_agent = TargetAgent(config)
+
+# or get target agent if it already exists
+# target_agent = TargetAgent.get(identifier="id_1234567")
+
+# Create Maihem Agent
+maihem_agent = MaihemAgent(
+    name="Maihem_Agent_1", # Name of the Maihem Agent, does not have to be unique
+    metric_to_eval="helpfulness", # (Optional) Select metric to evaluate, default is "helpfulness"
+    behavior_prompt="Vent about stress in the workplace.", # (Optional) Guide the behavior of the Maihem Agent with a prompt describing the workflow of the user
+    # behavior_workflow={}, # (Optional) Provide a dictionary of the workflow to guide the behavior of the Maihem Agent, or use our tool to help you generate this
+    # add_personality=True, # (Optional) Add personality to the agent
+    turns_max=8, # (Optional) Set maximum number of turns, default is 10
+)
+
+# (Optional): Connect Maihem Agent to Target Agent
+maihem_agent.connect_target_agent(target_agent)
+
+# (Optional): Define static responses and number of static turns
+maihem_agent.set_static_message_list(
+    static_responses=["I'm here to help you.", "How can I support you today?"], # Optional, list of static responses, dynamic after
+)
+
+maihem_agent.set_static_n_turns(
+    n=2, # Number of static responses to use from previous conversation generated by the Maihem Agent, after which dynamic responses are generated
+)
+
+# (Optional): Add documents for RAG evaluations
+# Load documents that the Maihem Agent can use to generate questions for the Target Agent
+maihem_agent.add_documents(paths=["./data/journaling.txt", "./data/habit-building.txt", "./data/emotional-support.txt"])
+
+# Chat with the Maihem Agent
+response = maihem_agent.chat(
+    target_agent_message="Hi, how can I help you today?", # Message from the Target Agent
+    contexts=["<Context 1>", "<Context 2>"], # (Optional) Pass list of contexts for RAG evaluations
+)
+
+# Get message content and evals
+response.content
+response.eval
+
+# or pass callable chat function to automate entire conversation
+target_agent.set_chat_function(chat_function)
+maihem_agent.connect_target_agent(target_agent)
+maihem_agent.generate_conversation()
+
+# Get conversation and evals
+maihem_agent.conversation
+maihem_agent.evals
\ No newline at end of file
diff --git a/src/maihem/api.py b/src/maihem/api.py
index bead931..3c15275 100644
--- a/src/maihem/api.py
+++ b/src/maihem/api.py
@@ -70,6 +70,7 @@
 
 from maihem.errors import handle_http_errors, ErrorResponse
 from maihem.schemas.tests import TestStatusEnum
+from maihem.logger import get_logger
 
 
 class MaihemHTTPClientBase:
@@ -161,7 +162,6 @@ def upsert_test(self, req: APISchemaTestCreateRequest) -> APISchemaTest:
                 x_api_key=self.token,
                 body=req,
             )
-
         if response.status_code != 201 and response.status_code != 200:
             error_dict = json.loads(response.content)
             handle_http_errors(error_resp=ErrorResponse.from_dict(error_dict))
diff --git a/src/maihem/api_client/maihem_client/api/agents/agents_upsert_agent_target.py b/src/maihem/api_client/maihem_client/api/agents/agents_upsert_agent_target.py
index 0cfbe0f..d3c193a 100644
--- a/src/maihem/api_client/maihem_client/api/agents/agents_upsert_agent_target.py
+++ b/src/maihem/api_client/maihem_client/api/agents/agents_upsert_agent_target.py
@@ -1,5 +1,5 @@
 from http import HTTPStatus
-from typing import Any, Dict, Optional, Union, cast
+from typing import Any, Dict, Optional, Union
 
 import httpx
 
@@ -36,7 +36,7 @@ def _get_kwargs(
 
 def _parse_response(
     *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Optional[Union[APISchemaAgentTarget, Any, ErrorResponse]]:
+) -> Optional[Union[APISchemaAgentTarget, ErrorResponse]]:
     if response.status_code == HTTPStatus.CREATED:
         response_201 = APISchemaAgentTarget.from_dict(response.json())
 
@@ -62,7 +62,8 @@ def _parse_response(
 
         return response_504
     if response.status_code == HTTPStatus.OK:
-        response_200 = cast(Any, None)
+        response_200 = APISchemaAgentTarget.from_dict(response.json())
+
         return response_200
     if client.raise_on_unexpected_status:
         raise errors.UnexpectedStatus(response.status_code, response.content)
@@ -72,7 +73,7 @@ def _parse_response(
 
 def _build_response(
     *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Response[Union[APISchemaAgentTarget, Any, ErrorResponse]]:
+) -> Response[Union[APISchemaAgentTarget, ErrorResponse]]:
     return Response(
         status_code=HTTPStatus(response.status_code),
         content=response.content,
@@ -86,7 +87,7 @@ def sync_detailed(
     client: AuthenticatedClient,
     body: APISchemaAgentTargetCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Response[Union[APISchemaAgentTarget, Any, ErrorResponse]]:
+) -> Response[Union[APISchemaAgentTarget, ErrorResponse]]:
     """Upsert target agent by identifier
 
      Create or update a target agent by identifier
@@ -100,7 +101,7 @@ def sync_detailed(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Response[Union[APISchemaAgentTarget, Any, ErrorResponse]]
+        Response[Union[APISchemaAgentTarget, ErrorResponse]]
     """
 
     kwargs = _get_kwargs(
@@ -120,7 +121,7 @@ def sync(
     client: AuthenticatedClient,
     body: APISchemaAgentTargetCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Optional[Union[APISchemaAgentTarget, Any, ErrorResponse]]:
+) -> Optional[Union[APISchemaAgentTarget, ErrorResponse]]:
     """Upsert target agent by identifier
 
      Create or update a target agent by identifier
@@ -134,7 +135,7 @@ def sync(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Union[APISchemaAgentTarget, Any, ErrorResponse]
+        Union[APISchemaAgentTarget, ErrorResponse]
     """
 
     return sync_detailed(
@@ -149,7 +150,7 @@ async def asyncio_detailed(
     client: AuthenticatedClient,
     body: APISchemaAgentTargetCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Response[Union[APISchemaAgentTarget, Any, ErrorResponse]]:
+) -> Response[Union[APISchemaAgentTarget, ErrorResponse]]:
     """Upsert target agent by identifier
 
      Create or update a target agent by identifier
@@ -163,7 +164,7 @@ async def asyncio_detailed(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Response[Union[APISchemaAgentTarget, Any, ErrorResponse]]
+        Response[Union[APISchemaAgentTarget, ErrorResponse]]
     """
 
     kwargs = _get_kwargs(
@@ -181,7 +182,7 @@ async def asyncio(
     client: AuthenticatedClient,
     body: APISchemaAgentTargetCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Optional[Union[APISchemaAgentTarget, Any, ErrorResponse]]:
+) -> Optional[Union[APISchemaAgentTarget, ErrorResponse]]:
     """Upsert target agent by identifier
 
      Create or update a target agent by identifier
@@ -195,7 +196,7 @@ async def asyncio(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Union[APISchemaAgentTarget, Any, ErrorResponse]
+        Union[APISchemaAgentTarget, ErrorResponse]
     """
 
     return (
diff --git a/src/maihem/clients.py b/src/maihem/clients.py
index cb71932..8ecaf3e 100644
--- a/src/maihem/clients.py
+++ b/src/maihem/clients.py
@@ -135,6 +135,42 @@ def create_target_agent(
 
         logger.info(f"Successfully created target agent {identifier}!")
         return agent_target
+    
+    def upsert_target_agent(
+        self,
+        identifier: str,
+        role: str,
+        industry: str,
+        description: str,
+        name: Optional[str] = None,
+        language: Optional[str] = "en",
+    ) -> TargetAgent:
+        logger = get_logger()
+        logger.info(f"Creating target agent {identifier}...")
+        resp = None
+        try:
+            resp = self._maihem_api_client.upsert_agent_target(
+                req=APISchemaAgentTargetCreateRequest(
+                    identifier=identifier,
+                    name=name,
+                    role=role,
+                    industry=industry,
+                    description=description,
+                    language=language
+                )
+            )
+        except errors.ErrorBase as e:
+            errors.handle_base_error(e)
+
+        agent_target = None
+
+        try:
+            agent_target = TargetAgent.model_validate(resp.to_dict())
+        except ValidationError as e:
+            errors.handle_schema_validation_error(e)
+
+        logger.info(f"Successfully created target agent {identifier}!")
+        return agent_target
 
     def get_target_agent(self, identifier: str) -> TargetAgent:
         resp = None
@@ -201,6 +237,54 @@ def create_test(
 
         logger.info(f"Successfull created test {identifier}!")
         return test
+    
+    def upsert_test(
+        self,
+        identifier: str,
+        initiating_agent: AgentType = AgentType.MAIHEM,
+        name: Optional[str] = None,
+        maihem_agent_behavior_prompt: Optional[str] = None,
+        conversation_turns_max: Optional[int] = 10,
+        metrics_config: Optional[Dict] = None,
+    ) -> Test:
+        resp = None
+        logger = get_logger()
+        logger.info(f"Creating test {identifier}...")
+
+        if isinstance(initiating_agent, str):
+            try:
+                initiating_agent = AgentType[initiating_agent.upper()]
+            except KeyError as e:
+                raise errors.raise_request_validation_error(
+                    f"Invalid agent type: {initiating_agent}"
+                ) from e
+
+        try:
+            metrics_config = APISchemaTestCreateRequestMetricsConfig.from_dict(
+                metrics_config
+            )
+            resp = self._maihem_api_client.upsert_test(
+                req=APISchemaTestCreateRequest(
+                    identifier=identifier,
+                    name=name,
+                    initiating_agent=initiating_agent,
+                    conversation_turns_max=conversation_turns_max,
+                    agent_maihem_behavior_prompt=maihem_agent_behavior_prompt,
+                    metrics_config=metrics_config,
+                    is_dev_mode=True
+                )
+            )
+        except errors.ErrorBase as e:
+            errors.handle_base_error(e)
+
+        test = None
+
+        try:
+            test = Test.model_validate(resp.to_dict())
+        except ValidationError as e:
+            errors.handle_schema_validation_error(e)
+
+        return test
 
     def get_test(self, identifier: str) -> Test:
         resp = None
@@ -320,6 +404,100 @@ def create_test_run(
             self._maihem_api_client.update_test_run_status(
                 test_run_id=test_run.id, status=TestStatusEnum.FAILED
             )
+            
+    def create_test_run_dev_mode(
+        self,
+        test_identifier: str,
+        target_agent: TargetAgent,
+        concurrent_conversations: int = 1,
+    ) -> TestRun:
+        try:
+            if not target_agent._chat_function:
+                errors.raise_request_validation_error(
+                    "Chat function the target agent must be passed."
+                )
+
+            test = self._maihem_api_client.get_test_by_identifier(test_identifier)
+
+            resp = None
+
+            logger = get_logger()
+
+            try:
+                with yaspin(
+                    Spinners.arc,
+                    text="Creating Maihem Agent, this might take a minute...",
+                ) as sp:
+                    resp = self._maihem_api_client.create_test_run(
+                        test_id=test.id,
+                        agent_target_id=target_agent.id,
+                    )
+            except errors.ErrorBase as e:
+                errors.handle_base_error(e)
+
+            test_run = None
+
+            try:
+                test_run = TestRun.model_validate(resp.to_dict())
+            except ValidationError as e:
+                errors.handle_schema_validation_error(e)
+
+            test_run_conversations = self.get_test_run_conversations(test_run.id)
+            conversation_ids = test_run_conversations.conversation_ids
+
+            # logger.info(
+            #     f"Test run results (UI): {self._base_url_ui}/evaluations/{test_run.id}"
+            # )
+
+            # with tqdm(
+            #     len(conversation_ids),
+            #     total=len(conversation_ids),
+            #     unit="conversation",
+            #     colour="green",
+            #     desc=f"Test run ({test_run.id})",
+            #     position=0,
+            # ) as progress:
+            with ThreadPoolExecutor(
+                max_workers=concurrent_conversations
+            ) as executor:
+                future_to_conversation_id = {
+                    executor.submit(
+                        self._run_conversation,
+                        test_run.id,
+                        conversation_id,
+                        test,
+                        target_agent,
+                        progress_bar_position=i + 1,
+                    ): conversation_id
+                    for i, conversation_id in enumerate(conversation_ids)
+                }
+
+                for future in as_completed(future_to_conversation_id):
+                    conversation_id = future_to_conversation_id[future]
+                    try:
+                        future.result()
+                    except errors.ErrorBase as e:
+                        logger.error(
+                            f"Error running conversation ({conversation_id}): {e.message}"
+                        )
+                        # progress.colour = "red"
+                    # finally:
+                        # progress.update()
+
+            print("\n" + "-" * 50 + "\n")
+            logger.info(f"Conversation ({test_run.id}) completed!")
+            return test_run
+        except KeyboardInterrupt:
+            logger.info("Keyboard interrupt detected. Canceling test...")
+            self._maihem_api_client.update_test_run_status(
+                test_run_id=test_run.id, status=TestStatusEnum.CANCELED
+            )
+            logger.info("Test run canceled!")
+        except Exception as e:
+            logger.error(f"Error: {e}. Ending test...")
+            self._maihem_api_client.update_test_run_status(
+                test_run_id=test_run.id, status=TestStatusEnum.FAILED
+            )
 
     def get_test_run_conversations(self, test_run_id: str) -> TestRunConversations:
         resp = None
diff --git a/src/maihem/errors.py b/src/maihem/errors.py
index 74e4288..881059d 100644
--- a/src/maihem/errors.py
+++ b/src/maihem/errors.py
@@ -14,6 +14,7 @@
 
 class ErrorCodes(str, Enum):
     ERR_CHAT_FUNCTION = "err_chat_function"
+    ERR_CONFIG_FILE = "err_config_file"
     ERR_INTERNAL_SERVER = ErrorCodesAPI.ERR_INTERNAL_SERVER
     ERR_NOT_FOUND = ErrorCodesAPI.ERR_NOT_FOUND
     ERR_REQUEST_VALIDATION = ErrorCodesAPI.ERR_REQUEST_VALIDATION
@@ -65,6 +66,10 @@ def __init__(self, error_resp: ErrorResponse) -> None:
 class ChatFunctionError(ErrorBase):
     def __init__(self, error_resp: ErrorResponse) -> None:
         super().__init__(error_resp)
+        
+class ConfigFileError(ErrorBase):
+    def __init__(self, error_resp: ErrorResponse) -> None:
+        super().__init__(error_resp)
 
 
 def handle_http_errors(error_resp: ErrorResponse):
@@ -135,3 +140,14 @@ def raise_chat_function_error(message: str):
             request_id=None,
         )
     )
+    
+    
+def raise_config_file_error(message: str):
+    raise ConfigFileError(
+        ErrorResponse(
+            error=ErrorResponseError(
+                code=ErrorCodes.ERR_CONFIG_FILE, message=message, detail=None
+            ),
+            request_id=None,
+        )
+    )
diff --git a/src/maihem/schemas/agents.py b/src/maihem/schemas/agents.py
index d1a77f0..fc7eda1 100644
--- a/src/maihem/schemas/agents.py
+++ b/src/maihem/schemas/agents.py
@@ -1,14 +1,21 @@
 from datetime import datetime
 from pydantic import BaseModel
+
+from enum import Enum
+import os
+from pydantic_extra_types.language_code import LanguageAlpha2
 from typing import Callable, Optional, Tuple, List, Dict
+
 import maihem.errors as errors
-from enum import Enum
 from maihem.logger import get_logger
-from pydantic_extra_types.language_code import LanguageAlpha2
-import os
 from maihem.utils import extract_text
 
 
+class AgentType(str, Enum):
+    MAIHEM = "maihem"
+    TARGET = "target"
+
+
 class TargetAgent(BaseModel):
     id: str
     created_at: datetime
@@ -87,8 +94,3 @@ def _send_message(
             errors.raise_chat_function_error("Target agent chat function not set")
         response, contexts = self._chat_function(conversation_id, message)
         return response, contexts
-
-
-class AgentType(str, Enum):
-    MAIHEM = "maihem"
-    TARGET = "target"
diff --git a/src/maihem/simulator.py b/src/maihem/simulator.py
new file mode 100644
index 0000000..5099742
--- /dev/null
+++ b/src/maihem/simulator.py
@@ -0,0 +1,101 @@
+from typing import List, Callable, Dict
+import yaml
+
+from maihem.clients import Maihem
+from maihem.errors import raise_config_file_error
+from maihem.logger import get_logger
+from maihem.schemas.tests import SimulatedConversation
+
+
+class Simulator:
+    """
+    Class to simulate conversations between the Maihem Agent and the Target Agent
+    """
+    
+    @classmethod
+    def conversation(
+        cls, 
+        chat_function: Callable,
+        target_agent_identifier: str,
+        maihem_agent_identifier: str,      
+        config_path: str = "./config.yaml",    
+    ) -> SimulatedConversation:
+        """Simulate a single conversation between the Maihem Agent and the Target Agent"""
+        logger = get_logger()
+        
+        # Read config.yaml file
+        config = cls._load_config(config_path, target_agent_identifier, maihem_agent_identifier)
+        
+        maihem_client = Maihem()
+        
+        logger.info("Initiating conversation...")
+        target_agent = maihem_client.upsert_target_agent(
+            identifier=target_agent_identifier, 
+            role=config["target_agent"][target_agent_identifier]["role"],
+            industry=config["target_agent"][target_agent_identifier]["industry"],
+            description=config["target_agent"][target_agent_identifier]["description"],
+            name=config["target_agent"][target_agent_identifier]["name"],
+            language=config["target_agent"][target_agent_identifier]["language"],
+        )
+        target_agent.set_chat_function(chat_function)
+        
+        metrics_config = {config["metric"]: 1}
+        
+        test = maihem_client.upsert_test(
+            identifier=maihem_agent_identifier,
+            initiating_agent=config["initiating_agent"],
+            name=config["maihem_agent"][maihem_agent_identifier]["name"],
+            maihem_agent_behavior_prompt=config["maihem_agent"][maihem_agent_identifier]["behavior_prompt"],
+            conversation_turns_max=config["maihem_agent"][maihem_agent_identifier]["conversation_turns_max"],
+            metrics_config=metrics_config
+        )
+        
+        test_run = maihem_client.create_test_run_dev_mode(
+            test_identifier=test.identifier,
+            target_agent=target_agent
+        )
+        
+        conversation = maihem_client.get_test_run_result_conversations(test_run.id)
+        
+        return SimulatedConversation(conversation)
+    
+    @classmethod
+    def _load_config(
+        cls, 
+        config_path: str, 
+        target_agent_identifier: str, 
+        maihem_agent_identifier: str
+    ) -> Dict:
+        logger = get_logger()
+        try:
+            with open(config_path, 'r') as file:
+                config = yaml.safe_load(file)
+        except Exception as e:
+            logger.error(f"Error reading config file: {str(e)}")    
+            raise_config_file_error(f"Error reading config file: {str(e)}")
+            
+        try:
+            assert "metric" in config, "'metric' key not found in config file"
+            assert "initiating_agent" in config, "'initiating_agent' key not found in maihem_agent"
+            assert "target_agent" in config, "'target_agent' key not found in config file"
+            assert "maihem_agent" in config, "'maihem_agent' key not found in config file" 
+            assert target_agent_identifier in config["target_agent"].keys(), f"Target agent '{target_agent_identifier}' not found in config file"
+            assert maihem_agent_identifier in config["maihem_agent"].keys(), f"Maihem agent '{maihem_agent_identifier}' not found in config file"
+            
+            for identifier, target_agent in config["target_agent"].items():
+                assert "role" in target_agent, f"'role' key not found in target agent {identifier}"
+                assert "industry" in target_agent, f"'industry' key not found in target agent {identifier}"
+                assert "description" in target_agent, f"'description' key not found in target agent {identifier}"
+                target_agent["name"] = target_agent.get("name")
+                target_agent["language"] = target_agent.get("language")
+                
+            for identifier, maihem_agent in config["maihem_agent"].items():
+                maihem_agent["name"] = maihem_agent.get("name")
+                maihem_agent["behavior_prompt"] = maihem_agent.get("behavior_prompt")
+                maihem_agent["conversation_turns_max"] = maihem_agent.get("conversation_turns_max")
+                       
+        except Exception as e:
+            logger.error(f"Error reading config file: {str(e)}")    
+            raise_config_file_error(f"Error reading config file: {str(e)}")
+            
+        return config
\ No newline at end of file
diff --git a/src/run_maihem.py b/src/run_maihem.py
new file mode 100644
index 0000000..089e3bf
--- /dev/null
+++ b/src/run_maihem.py
@@ -0,0 +1,57 @@
+import argparse
+import os
+from typing import List, Tuple
+from maihem.simulator import Simulator
+
+# Set the Maihem API key as an environment variable
+# os.environ['MAIHEM_API_KEY'] = '<YOUR_API_KEY>'
+
+parser = argparse.ArgumentParser()
+
+# Choose identifier of the target agent, to be defined in config.yaml
+parser.add_argument("--target_agent_identifier", default="tg_ag_1_upsert_1", type=str)
+
+# Choose identifier of the maihem agent, to be defined in config.yaml
+parser.add_argument("--maihem_agent_identifier", default="mh_ag_1_upsert_1", type=str)
+
+# Path to config.yaml
+parser.add_argument("--config_path", default="./src/config.yaml", type=str)
+
+
+# Example of chat function
+def chat_function(conversation_id: str, agent_maihem_message: str) -> Tuple[str, List[str]]:
+    """
+    Callable chat function to wrap your target agent to be tested.
+
+    params:
+    - conversation_id: unique conversation ID generated by Maihem, use it to keep track of different conversations
+    - agent_maihem_msg: message from the Maihem Agent
+
+    return:
+    - agent_target_message [str] - message from your Target Agent in response to the Maihem Agent message
+    - contexts [List[str]] - list of contexts (required for RAG evaluations), pass empty list if not needed
+    """
+    
+    # The message from your Target Agent. 
+    # Connect your agent here to respond to agent_maihem_message. 
+    # Keep track of same or new conversations with conversation_id
+    agent_target_message = "Hi, how can I help you today?"
+    
+    # List of contexts for RAG evaluations, pass empty list if not needed
+    contexts = [] 
+    # contexts = ["<Context 1>", "<Context 2>"] 
+    
+    return agent_target_message, contexts
+
+
+if __name__ == "__main__":
+    
+    conversation = Simulator.conversation(
+        chat_function, 
+        parser.parse_args().target_agent_identifier, 
+        parser.parse_args().maihem_agent_identifier, 
+        parser.parse_args().config_path
+    )
+    
+    print(conversation.messages)
+    print(conversation.evaluation)


Test patch:
diff --git a/design_maihem_test.py b/design_maihem_test.py
new file mode 100644
index 0000000..e699489
--- /dev/null
+++ b/design_maihem_test.py
@@ -0,0 +1,73 @@
+from typing import List, Tuple
+
+from maihem.clients import Maihem
+
+# Add API key to MAIHEM_API_KEY environment variable
+# os.environ["MAIHEM_API_KEY"] = api_key
+maihem_client = Maihem()
+
+target_agent = maihem_client.create_target_agent(
+    name="PG_Chatbot_1", # Name of the Target Agent, does not have to be unique
+    identifier="tg_ag_1234567", # (Optional) Define unique identifier for your Target Agent, or we create one for you
+    target_agent_role="mental health assistant",
+    target_agent_description="A chatbot for emotional support"
+)
+
+# or get target agent if it already exists
+target_agent = TargetAgent.get(identifier="id_1234567")
+
+# Example of chat function to call your Target Agent
+def chat_function(conversation_id: str, agent_maihem_message: str) -> Tuple[str, List[str]]:
+    """
+    Callable chat function to wrap your target agent to be tested.
+
+    params:
+    - conversation_id: unique conversation ID generated by Maihem, use it to keep track of different conversations
+    - agent_maihem_msg: message from the Maihem Agent
+
+    return:
+    - agent_target_message [str] - message from your Target Agent in response to the Maihem Agent message
+    - contexts [List[str]] - list of contexts (required for RAG evaluations), pass empty list if not needed
+    """
+    
+    # The message from your Target Agent. 
+    # Connect your agent here to respond to agent_maihem_message. 
+    # Keep track of same or new conversations with conversation_id
+    agent_target_message = "Hi, how can I help you today?"
+    
+    # List of contexts for RAG evaluations, pass empty list if not needed
+    contexts = ["<Context 1>", "<Context 2>"] 
+    
+    return agent_target_message, contexts
+    
+# Set the chat function to the Target Agent, so we can automatically chat with your Target Agent
+target_agent.set_chat_function(chat_function)
+
+# For RAG testing, add a list of documents so Maihem can generate questions from them to evaluate your Target Agent
+target_agent.add_documents(documents=["./test_data/RAG_PG_essays.pdf"])
+
+# Create a Test
+test = MaihemClient.create_test(
+    name="Test PG Essays", # Name of the Test, does not have to be unique
+    identifier="test_pg_essays_1234", # (Optional) Define unique identifier for the Test, or we create one for you
+    eval="helpfulness", # (Optional) Select metric to evaluate, default is "helpfulness"
+    workflow_prompt="Vent about stress in the workplace.", # (Optional) Guide the behavior of the Maihem Agent with a prompt describing the workflow of the user
+    workflow_dict={}, # (Optional) Provide a dictionary of the workflow to guide the behavior of the Maihem Agent, or use our tool to help you generate this
+    add_personality=True, # (Optional) Add personality to the agent
+    turns_max=10, # (Optional) Set maximum number of turns
+)
+
+# # Create a Test Run from a Test and a Target Agent
+test_run = MaihemClient.create_test_run(
+    test_identifier="test_pg_essays", # The identifier of the Test you want to run
+    target_agent=target_agent, # The Target Agent you want to test
+    concurrent_conversations=10, # How many conversations to run concurrently
+)
+
+# Get the results of the Test Run (status and how many conversations passed or failed)
+results = MaihemClient.get_test_run_result(test_run_id=test_run.id)
+print(results)
+
+# Get the conversations generated in the Test Run
+conversations = MaihemClient.get_test_run_result_conversations(test_run_id=test_run.id)
+print(conversations)
\ No newline at end of file
diff --git a/src/maihem/api_client/maihem_client/api/tests/tests_upsert_test.py b/src/maihem/api_client/maihem_client/api/tests/tests_upsert_test.py
index 522cf12..8f80a70 100644
--- a/src/maihem/api_client/maihem_client/api/tests/tests_upsert_test.py
+++ b/src/maihem/api_client/maihem_client/api/tests/tests_upsert_test.py
@@ -1,5 +1,5 @@
 from http import HTTPStatus
-from typing import Any, Dict, Optional, Union, cast
+from typing import Any, Dict, Optional, Union
 
 import httpx
 
@@ -36,7 +36,7 @@ def _get_kwargs(
 
 def _parse_response(
     *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Optional[Union[APISchemaTest, Any, ErrorResponse]]:
+) -> Optional[Union[APISchemaTest, ErrorResponse]]:
     if response.status_code == HTTPStatus.CREATED:
         response_201 = APISchemaTest.from_dict(response.json())
 
@@ -62,7 +62,8 @@ def _parse_response(
 
         return response_504
     if response.status_code == HTTPStatus.OK:
-        response_200 = cast(Any, None)
+        response_200 = APISchemaTest.from_dict(response.json())
+
         return response_200
     if client.raise_on_unexpected_status:
         raise errors.UnexpectedStatus(response.status_code, response.content)
@@ -72,7 +73,7 @@ def _parse_response(
 
 def _build_response(
     *, client: Union[AuthenticatedClient, Client], response: httpx.Response
-) -> Response[Union[APISchemaTest, Any, ErrorResponse]]:
+) -> Response[Union[APISchemaTest, ErrorResponse]]:
     return Response(
         status_code=HTTPStatus(response.status_code),
         content=response.content,
@@ -86,7 +87,7 @@ def sync_detailed(
     client: AuthenticatedClient,
     body: APISchemaTestCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Response[Union[APISchemaTest, Any, ErrorResponse]]:
+) -> Response[Union[APISchemaTest, ErrorResponse]]:
     """Upsert test by identifier
 
      Create or update a test by identifier
@@ -100,7 +101,7 @@ def sync_detailed(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Response[Union[APISchemaTest, Any, ErrorResponse]]
+        Response[Union[APISchemaTest, ErrorResponse]]
     """
 
     kwargs = _get_kwargs(
@@ -120,7 +121,7 @@ def sync(
     client: AuthenticatedClient,
     body: APISchemaTestCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Optional[Union[APISchemaTest, Any, ErrorResponse]]:
+) -> Optional[Union[APISchemaTest, ErrorResponse]]:
     """Upsert test by identifier
 
      Create or update a test by identifier
@@ -134,7 +135,7 @@ def sync(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Union[APISchemaTest, Any, ErrorResponse]
+        Union[APISchemaTest, ErrorResponse]
     """
 
     return sync_detailed(
@@ -149,7 +150,7 @@ async def asyncio_detailed(
     client: AuthenticatedClient,
     body: APISchemaTestCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Response[Union[APISchemaTest, Any, ErrorResponse]]:
+) -> Response[Union[APISchemaTest, ErrorResponse]]:
     """Upsert test by identifier
 
      Create or update a test by identifier
@@ -163,7 +164,7 @@ async def asyncio_detailed(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Response[Union[APISchemaTest, Any, ErrorResponse]]
+        Response[Union[APISchemaTest, ErrorResponse]]
     """
 
     kwargs = _get_kwargs(
@@ -181,7 +182,7 @@ async def asyncio(
     client: AuthenticatedClient,
     body: APISchemaTestCreateRequest,
     x_api_key: Union[None, Unset, str] = UNSET,
-) -> Optional[Union[APISchemaTest, Any, ErrorResponse]]:
+) -> Optional[Union[APISchemaTest, ErrorResponse]]:
     """Upsert test by identifier
 
      Create or update a test by identifier
@@ -195,7 +196,7 @@ async def asyncio(
         httpx.TimeoutException: If the request takes longer than Client.timeout.
 
     Returns:
-        Union[APISchemaTest, Any, ErrorResponse]
+        Union[APISchemaTest, ErrorResponse]
     """
 
     return (
diff --git a/src/maihem/schemas/tests.py b/src/maihem/schemas/tests.py
index 0ec15aa..192f101 100644
--- a/src/maihem/schemas/tests.py
+++ b/src/maihem/schemas/tests.py
@@ -123,3 +123,23 @@ def to_dict(self):
 
     class Config:
         arbitrary_types_allowed = True
+        
+        
+class SimulatedConversation:
+    
+    def __init__(self, conversation: TestRunResultConversations):
+        self.messages = self._convert_conv_to_message_list(conversation)
+        self.evaluation = conversation.conversations[0].evaluations[0].explanation
+    
+    def _convert_conv_to_message_list(self, conversation: TestRunResultConversations):
+        conversation.conversations[0].conversation_turns[2].conversation_messages[0].agent_type
+        conv = conversation.conversations[0]
+        message_list = []
+        
+        for turn in conv.conversation_turns:
+            for message in turn.conversation_messages:
+                message_list.append({
+                    message.agent_type.value: message.content
+                })
+                
+        return message_list