From 05243ff4336a1ec64b4a7ce0a571dacde22ba8d3 Mon Sep 17 00:00:00 2001 From: Robert Jambrecic Date: Wed, 3 Jul 2024 14:53:31 +0200 Subject: [PATCH] Change raw SQL queries to prisma queries --- captn/captn_agents/application.py | 5 +++-- captn/captn_agents/db_queries.py | 22 ++++++------------- openai_agent/application.py | 5 +++-- .../ci/captn/captn_agents/test_application.py | 12 +++++++++- .../test_openai_agent_application.py | 11 ++++++++-- 5 files changed, 33 insertions(+), 22 deletions(-) diff --git a/captn/captn_agents/application.py b/captn/captn_agents/application.py index 82244d0c..1e41e88a 100644 --- a/captn/captn_agents/application.py +++ b/captn/captn_agents/application.py @@ -7,6 +7,7 @@ import httpx import openai import pandas as pd +import prisma from autogen.io.websockets import IOWebsockets from fastapi import APIRouter, File, Form, HTTPException, Query, UploadFile from prometheus_client import Counter @@ -240,7 +241,7 @@ async def get_smart_suggestions( user_id: Annotated[int, Query(description="The user id")], ) -> List[str]: user_initial_team = await get_initial_team(user_id) - if isinstance(user_initial_team, dict): - return user_initial_team["smart_suggestions"] # type: ignore[no-any-return] + if isinstance(user_initial_team, prisma.models.UserInitialTeam): + return user_initial_team.initial_team.smart_suggestions # type: ignore[no-any-return] return DEFAULT_SMART_SUGGESTIONS diff --git a/captn/captn_agents/db_queries.py b/captn/captn_agents/db_queries.py index 780fc51e..e9e5a416 100644 --- a/captn/captn_agents/db_queries.py +++ b/captn/captn_agents/db_queries.py @@ -1,22 +1,14 @@ -from typing import Any, Union +from typing import Optional, Union + +from prisma.models import UserInitialTeam from .helpers import get_db_connection -async def get_initial_team(user_id: Union[int, str]) -> Any: +async def get_initial_team(user_id: Union[int, str]) -> Optional[UserInitialTeam]: async with get_db_connection() as db: - query = f"""SELECT - uit.id AS user_initial_team_id, - uit.user_id, - it.id AS initial_team_id, - it.name AS initial_team_name, - it.smart_suggestions -FROM - "UserInitialTeam" uit -JOIN - "InitialTeam" it ON uit.initial_team_id = it.id -WHERE - uit.user_id = {int(user_id)}""" # nosec: [B608] - user_initial_team = await db.query_first(query) + user_initial_team = await db.userinitialteam.find_first( + where={"user_id": int(user_id)}, include={"initial_team": True} + ) return user_initial_team diff --git a/openai_agent/application.py b/openai_agent/application.py index 3d1cde12..58fbacfb 100644 --- a/openai_agent/application.py +++ b/openai_agent/application.py @@ -2,6 +2,7 @@ from os import environ from typing import Dict, List, Optional, Union +import prisma from fastapi import APIRouter, BackgroundTasks from openai import AsyncAzureOpenAI from pydantic import BaseModel @@ -273,10 +274,10 @@ async def chat( request: AzureOpenAIRequest, background_tasks: BackgroundTasks ) -> Dict[str, Union[Optional[str], int, Union[str, Optional[SmartSuggestions]]]]: user_initial_team = await get_initial_team(request.user_id) - if isinstance(user_initial_team, dict): + if isinstance(user_initial_team, prisma.models.UserInitialTeam): return { "team_status": "inprogress", - "team_name": user_initial_team["initial_team_name"], + "team_name": user_initial_team.initial_team.name, "team_id": request.chat_id, "customer_brief": "This is my customer brief.", "conversation_name": "Team of Experts", diff --git a/tests/ci/captn/captn_agents/test_application.py b/tests/ci/captn/captn_agents/test_application.py index 517cd72e..96fd535d 100644 --- a/tests/ci/captn/captn_agents/test_application.py +++ b/tests/ci/captn/captn_agents/test_application.py @@ -6,6 +6,7 @@ import autogen import pandas as pd +import prisma import pytest from autogen.io.websockets import IOWebsockets from fastapi import HTTPException @@ -513,11 +514,20 @@ def test_upload_csv_or_xlsx_file( class TestGetSmartSuggestions: + initial_team = prisma.models.InitialTeam( + id=1, + name="test_team", + smart_suggestions=["Boost sales", "Increase brand awareness"], + ) + user_initial_team = prisma.models.UserInitialTeam( + id=1, user_id=123, initial_team_id=1, initial_team=initial_team + ) + @pytest.mark.parametrize( ("return_value", "expected"), [ ( - {"smart_suggestions": ["Boost sales", "Increase brand awareness"]}, + user_initial_team, ["Boost sales", "Increase brand awareness"], ), (None, DEFAULT_SMART_SUGGESTIONS), diff --git a/tests/ci/openai_agent/test_openai_agent_application.py b/tests/ci/openai_agent/test_openai_agent_application.py index 80060ea9..8bb274b4 100644 --- a/tests/ci/openai_agent/test_openai_agent_application.py +++ b/tests/ci/openai_agent/test_openai_agent_application.py @@ -1,5 +1,6 @@ import unittest +import prisma import pytest from fastapi.testclient import TestClient @@ -30,8 +31,14 @@ def test_chat(self) -> None: with unittest.mock.patch( "openai_agent.application.get_initial_team" - ) as mock_get_initial_team: - mock_get_initial_team.return_value = {"initial_team_name": "test_team"} + ) as mock_get_user_initial_team: + initial_team = prisma.models.InitialTeam( + id=1, name="test_team", smart_suggestions=["test suggestion"] + ) + user_initial_team = prisma.models.UserInitialTeam( + id=1, user_id=123, initial_team_id=1, initial_team=initial_team + ) + mock_get_user_initial_team.return_value = user_initial_team response = client.post("/chat", json=request.model_dump())