random4.py

# Any import from src or main should be done inside setup_and_run function
# importing them here causes mock.patch to fail
import asyncio
import os
from typing import List, TypedDict
from unittest.mock import AsyncMock
from fastapi import testclient
import pytest
from unittest import mock


@mock.patch(
    "src.work.get_call_details",
)
@pytest.mark.asyncio
async def test_websocket_handler(
    mock_get_call_details,
):
    try:
        print("Running test")
        create_tenants_folder()
        mock_get_call_details.return_value = {
            "from_number": "7202537543",
            "to_number": "1234567890",
            "recording_url": "https://example.com/recording.mp3",
            "disconnect_reason": "completed",
            "start_time": 1703302407333,
            "end_time": 1703302428855,
        }

        anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")

        asana_project_id = os.environ.get("ASANA_PROJECT_ID")
        # If anthropic_api_key is not set, we mock the extract_info_from_transcript function
        if anthropic_api_key is None or anthropic_api_key == "":
            with mock.patch(
                "src.work.extract_info_from_transcript"
            ) as mock_extract_info_from_transcript:
                mock_extract_info(mock_extract_info_from_transcript)
                if asana_project_id is not None:
                    with mock.patch("src.agents.get_agent_data") as mock_get_agent_data:
                        mock_agent_data(mock_get_agent_data)
                        await setup_and_run()
                else:
                    await setup_and_run()
        else:
            if asana_project_id is not None:
                with mock.patch("src.agents.get_agent_data") as mock_get_agent_data:
                    mock_agent_data(mock_get_agent_data)
                    await setup_and_run()
            else:
                await setup_and_run()
    except Exception as e:
        print(e)
        assert False


async def setup_and_run():
    from src.llm_client import LlmClient
    from main import app, set_llm_client
    from src.utils.db import get_collection
    from src.work import get_asana_clients

    client = testclient.TestClient(app)
    llm_client = LlmClient()
    set_llm_client(llm_client)
    from src.llm_client import ResponseResponse, Utterance

    from secrets import token_hex

    print("Setting up test")

    call_id = token_hex(32)

    agent_id = "agent_fee6263535f3b5d6f1b5ad3f67"
    response_id = 0

    transcripts = []
    print("Connecting to websocket", call_id)
    dummy_transcript = get_dummy_transcript()

    last_sent_agent_message = 1
    last_sent_user_message = 0

    async def mock_generate_response(messages, response_id, agent_data):
        nonlocal last_sent_agent_message
        response = ResponseResponse(
            response_id=response_id,
            content=dummy_transcript[last_sent_agent_message]["content"],
            content_complete=True,
            end_call=False,
        )
        last_sent_agent_message += 2
        await llm_client.response_queue.put(response)

    llm_client._generate_response = AsyncMock(side_effect=mock_generate_response)
    calls_collection = get_collection("calls")

    with client.websocket_connect(f"/llm-websocket/{agent_id}/{call_id}") as websocket:
        while True:
            recv_data = websocket.receive_json()

            if recv_data["response_type"] == "config":
                websocket.send_json(
                    {
                        "interaction_type": "call_details",
                        "call": {
                            "call_type": "phone_call",
                            "from_number": "1234567890",
                            "to_number": "1234567890",
                            "direction": "inbound",
                            "call_id": call_id,
                            "agent_id": agent_id,
                            "call_status": "registered",
                            "retell_llm_dynamic_variables": {
                                "customer_name": "John Doe"
                            },
                            "opt_out_sensitive_data_storage": True,
                        },
                    }
                )

            if recv_data["response_type"] == "response":
                prev_response = ResponseResponse(**recv_data)

                if prev_response.response_id == 0:
                    call = calls_collection.find_one({"call_id": call_id})
                    assert call.get("call_id") == call_id
                    assert call.get("agent_id") == agent_id
                    assert len(call.get("transcript")) == 0

                if prev_response.end_call:
                    websocket.close()
                    break

                if response_id < 8:
                    response_id += 1
                    transcripts.append(
                        Utterance(role="agent", content=prev_response.content).dict()
                    )
                    transcripts.append(
                        Utterance(
                            role="user",
                            content=dummy_transcript[last_sent_user_message]["content"],
                        ).dict()
                    )

                    if response_id == 8:
                        transcripts.append(
                            Utterance(
                                role=dummy_transcript[-1]["role"],
                                content=dummy_transcript[-1]["content"],
                            ).dict()
                        )
                    last_sent_user_message += 2

                    payload = {
                        "interaction_type": "response_required",
                        "response_id": response_id,
                        "transcript": transcripts,
                    }

                    websocket.send_json(payload)
                else:
                    websocket.close()
                    break
        call = calls_collection.find_one({"call_id": call_id})

        count = 0

        while call.get("task_id") is None:
            call = calls_collection.find_one({"call_id": call_id})
            await asyncio.sleep(1)
            count += 1

            if count > 10:
                raise TimeoutError("Task not created")

        assert call.get("call_id") == call_id
        db_transcript = call.get("transcript")
        assert len(db_transcript) != 0

        assert len(db_transcript) == len(dummy_transcript) + 1

        assert db_transcript[0]["role"] == "agent"

        for i in range(1, len(dummy_transcript) - 1):
            assert db_transcript[i + 1]["role"] == dummy_transcript[i]["role"]
            assert db_transcript[i + 1]["content"] == dummy_transcript[i]["content"]

        tasks_api, _ = get_asana_clients(os.environ.get("ASANA_API_KEY"))
        created_task = tasks_api.get_task(
            call.get("task_id"),
            opts={
                "opt_fields": "name,notes,",
                "workspace": os.environ.get("ASANA_WORKSPACE_ID"),
            },
        )
        assert created_task["gid"] == call.get("task_id")
        assert created_task["name"] == "[AIPM] customer: John Doe"

        client.close()


class Message(TypedDict):
    role: str
    content: str


def create_tenants_folder():
    if not os.path.exists("./data/test_company"):
        os.makedirs("./data/test_company")
    tenants_data = """Property,Unit,Tenant,Status,Tenant Type,Phone Numbers,Emails,Move-in,Lease To,Rent,Deposit,Tenant Tags

"Chen - 942 - 942 28th St NE Auburn, WA 98002",Chen - 942,"Robert, Welch",Past,Other Occupant,Mobile: (843) 300-9925,,01/20/2022,01/31/2024,0.00,0.00,"RBP, MO Complete"
"Lund - 4306 40th Ave NE Tacoma, WA 98422",Lund,"Absie, Farhia M.",Current,Financially Responsible,Mobile: (206) 295-2381,[email protected],02/01/2023,02/28/2025,"2,950.00","2,750.00","RBP, RBP Enrollment Complete"
"Lund - 4306 40th Ave NE Tacoma, WA 98422",Lund,"Absie, Sasiya",Current,Other Occupant,,,02/01/2023,02/28/2025,"2,950.00","2,750.00","RBP, RBP Enrollment Complete"
"Nelson 1223 - 1223 S Verde St Tacoma, WA 98405",Nelson 1223,"Adamczyk, Michael C.",Current,Financially Responsible,Mobile: (720) 253-7543,[email protected],08/01/2023,07/31/2025,"2,500.00","2,800.00","RBP, RBP Enrollment Complete"
"""
    with open("./data/test_company/tenants.csv", "w") as f:
        f.write(tenants_data)


def mock_extract_info(mock_function):
    mock_function.return_value = {
        "caller_type": "customer",
        "name": "John Doe",
        "address": "123 Main St, Denver, CO 80202",
        "reason_for_call": "Maintenance",
    }


def mock_agent_data(mock_function):
    mock_function.return_value = {
        "agent_id": "agent_fee6263535f3b5d6f1b5ad3f67",
        "client_name": "Testing PMC Company",
        "asana_project_id": "1208699491890437",
        "data_folder": "./data/test_company",
        "instructions": """
Heating and air condition if hoter than 80f or colder than 50f
If windows are broken and it appears like burglary then call 911.
        """,
    }


def get_dummy_transcript() -> List[Message]:
    return [
        {
            "role": "user",
            "content": "I was told to call for a maintenance order that I put in.",
        },
        {
            "role": "agent",
            "content": "Thank you for reaching out. Can you please confirm your full name and the property address?",
        },
        {
            "role": "user",
            "content": "John Doe, 123 Main St, Denver, CO 80202",
        },
        {
            "role": "agent",
            "content": "Thank you, John. Could you please describe the issue you're experiencing with your unit?",
        },
        {
            "role": "user",
            "content": "It's a house. Um, my",
        },
        {
            "role": "agent",
            "content": "Thank you, John. Could you please describe the issue you're experiencing with your unit?",
        },
        {
            "role": "user",
            "content": "It's a house. Um, my heater's not going on.",
        },
        {
            "role": "agent",
            "content": "Thank you for letting me know. Given that it's a heating issue, could you please tell me the current temperature inside your house? This will help determine if it's an emergency situation.",
        },
        {
            "role": "user",
            "content": "I'm at work.",
        },
        {
            "role": "agent",
            "content": "understand. Since you're not currently at home, I'm unable to provide troubleshooting steps at this time. Please check the temperature when you return, and ",
        },
        {
            "role": "user",
            "content": "Can I talk to customer service?",
        },
        {
            "role": "agent",
            "content": "Certainly, John. I will pass your request to our customer service team. Can you please provide a contact number where they can reach you?",
        },
        {
            #
            "role": "user",
            "content": "7 02:07 5 5 4 03:02?",
        },
        {
            "role": "agent",
            "content": "understand the urgency, John. Since there's an infant at home and the heater isn't working, I'll create a work order for this emergency. A coordinator will contact you shortly to address the issue. If you have any further concerns, please let me know.",
        },
        {
            #
            "role": "user",
            "content": "Thank you.",
        },
        {
            "role": "agent",
            "content": "You're welcome, John. If you need anything else, feel free to reach out. Take care.",
        },
    ]