Merge pull request #62 from studio-recoding/dev

feat: case2 카테고리 기능 추가

app/database/chroma_db.py

@@ -10,7 +10,7 @@
 import os
 import datetime
 from dotenv import load_dotenv
-from app.dto.db_dto import AddScheduleDTO, RecommendationMainRequestDTO, ReportTagsRequestDTO
+from app.dto.db_dto import AddScheduleDTO, DeleteScheduleDTO, UpdateScheduleDTO, RecommendationMainRequestDTO, ReportTagsRequestDTO
 
 load_dotenv()
 CHROMA_DB_IP_ADDRESS = os.getenv("CHROMA_DB_IP_ADDRESS")
@@ -37,30 +37,76 @@ def check_db_heartbeat():
     chroma_client.heartbeat()
 
 # description: DB에서 검색하는 함수 - chat case 3에 사용
-async def search_db_query(query):
-    # 컬렉션 생성
-    # 컬렉션에 쿼리 전송
+async def search_db_query(member_id, query):
+    member = member_id
     result = schedules.query(
         query_texts=query,
-        n_results=5  # 결과에서 한 가지 문서만 반환하면 한강공원이, 두 가지 문서 반환하면 AI가 뜸->유사도가 이상하게 검사되는 것 같음
+        n_results=15,  # 결과에서 한 가지 문서만 반환하면 한강공원이, 두 가지 문서 반환하면 AI가 뜸->유사도가 이상하게 검사되는 것 같음
+        where={"member": {"$eq": int(member)}}
     )
     return result
 
 # description: DB에 저장하는 함수
 # 스프링 백엔드로부터 chroma DB에 저장할 데이터를 받아 DB에 추가한다.
 async def add_db_data(schedule_data: AddScheduleDTO):
+    schedule_date = schedule_data.schedule_datetime_start.split("T")[0]
+    year, month, date = map(int, schedule_date.split("-"))
+
+    schedules.add(
+        documents=[schedule_data.data],
+        ids=[str(schedule_data.schedule_id)],
+        metadatas=[{
+            "year": year,
+            "month": month,
+            "date": date,
+            "datetime_start": schedule_data.schedule_datetime_start,
+            "datetime_end": schedule_data.schedule_datetime_end,
+            "member": schedule_data.member_id,
+            "category": schedule_data.category,
+            "category_id": schedule_data.schedule_id,
+            "location": schedule_data.location,
+            "person": schedule_data.person
+        }]
+    )
+    return True
+
+
+# 메인페이지 한 줄 추천 기능에 사용하는 함수
+async def delete_db_data(schedule_data: DeleteScheduleDTO):
+    member_id = schedule_data.member_id
+    schedule_id = schedule_data.schedule_id
+    schedules.delete(
+        ids=[str(schedule_id)],
+        where={"member": {"$eq": int(member_id)}}
+    )
+    return True
+
+# 데이터베이스 업데이트 함수 정의
+async def update_db_data(schedule_data: UpdateScheduleDTO):
     schedule_date = schedule_data.schedule_datetime_start.split("T")[0]
     year = int(schedule_date.split("-")[0])
     month = int(schedule_date.split("-")[1])
     date = int(schedule_date.split("-")[2])
-    schedules.add(
+
+    # 기존 스케줄 업데이트 로직
+    schedules.update(
         documents=[schedule_data.data],
         ids=[str(schedule_data.schedule_id)],
-        metadatas=[{"year": year, "month": month, "date": date, "datetime_start": schedule_data.schedule_datetime_start, "datetime_end": schedule_data.schedule_datetime_end, "member": schedule_data.member_id, "category": schedule_data.category, "location": schedule_data.location, "person": schedule_data.person}]
+        metadatas=[{
+            "year": year,
+            "month": month,
+            "date": date,
+            "datetime_start": schedule_data.schedule_datetime_start,
+            "datetime_end": schedule_data.schedule_datetime_end,
+            "member": schedule_data.member_id,
+            "category": schedule_data.category,
+            "category_id": schedule_data.schedule_id,
+            "location": schedule_data.location,
+            "person": schedule_data.person
+        }]
     )
     return True
 
-# 메인페이지 한 줄 추천 기능에 사용하는 함수
 # 유저의 id, 해당 날짜로 필터링
 async def db_daily_schedule(user_data: RecommendationMainRequestDTO):
     member = user_data.member_id

app/database/connect_rds.py

@@ -0,0 +1,31 @@
+import os
+import pymysql
+from pymysql.cursors import DictCursor
+from dotenv import load_dotenv
+
+# .env 파일에서 환경 변수 로드
+load_dotenv()
+
+def get_rds_connection():
+    return pymysql.connect(
+        host=os.getenv('RDS_HOST'),
+        user=os.getenv('RDS_USER'),
+        password=os.getenv('RDS_PASSWORD'),
+        database=os.getenv('RDS_DATABASE'),
+        cursorclass=DictCursor
+    )
+
+def fetch_category_classification_data(member_id):
+    connection = get_rds_connection()
+    try:
+        with connection.cursor() as cursor:
+            sql = """
+            SELECT c.*
+            FROM category c
+            WHERE c.member_id = %s
+            """
+            cursor.execute(sql, (member_id,))
+            result = cursor.fetchall()
+            return result
+    finally:
+        connection.close()

app/dto/db_dto.py

@@ -7,7 +7,23 @@ class AddScheduleDTO(BaseModel):
     schedule_datetime_end: str
     schedule_id: int
     member_id: int
-    category: int
+    category: str
+    category_id: int
+    location: str
+    person: str
+
+class DeleteScheduleDTO(BaseModel):
+    schedule_id: int
+    member_id: int
+
+class UpdateScheduleDTO(BaseModel):
+    data: str
+    schedule_datetime_start: str
+    schedule_datetime_end: str
+    schedule_id: int
+    member_id: int
+    category: str
+    category_id: int
     location: str
     person: str
 

app/dto/openai_dto.py

@@ -2,8 +2,10 @@
 from typing import List
 
 class PromptRequest(BaseModel):
+    member_id: int
     prompt: str
     persona: str
+    chatType: str
 
 class ChatResponse(BaseModel):
     ness: str
@@ -23,19 +25,11 @@ class EmailResponse(BaseModel):
     text: str
     image: str
 
-<<<<<<< Updated upstream
 class ActivityDescription(BaseModel):
     activity: str
     imageTag: str
 
 class RecommendationResponse(BaseModel):
     ness: str
     activityList: List[ActivityDescription]
-=======
-class ActivityInfo(BaseModel):
-    activity: str
-    imageTag: str
-class RecommendationResponse(BaseModel):
-    ness: str
-    activityList: List[ActivityInfo]
->>>>>>> Stashed changes
+

app/prompt/openai_prompt.py

@@ -41,6 +41,8 @@ class Template:
                     Task: User Chat Classification
                     You are a case classifier integrated in scheduler application.
                     Please analyze User Chat according to the following criteria and return the appropriate case number (1, 2, 3).
+                    {chat_type}
+                
                     - Case 1: \
                     The question is a general information request, advice, or simple conversation, and does not require accessing the user's schedule database.
                     - Case 2: \
@@ -71,44 +73,60 @@ class Template:
                     User Chat: {question}
                     Answer:
                     """
+    chat_type_stt_template = """
+            You should keep in mind that this user's input was written using speech to text technology.
+            Therefore, there may be inaccuracies in the text due to errors in the STT process.
+            You need to consider this aspect when performing the given task.
+    """
+    chat_type_user_template = """
+    """
     case1_template = """
             {persona}
-            YOU MUST USE {output_language} TO RESPOND TO THE USER INPUT.
+            {chat_type}
+            YOU MUST USE {output_language} TO RESPOND TO THE USER INPUT. Current time is {current_time}. Respond to the user considering the current time.
             User input: {question}
             """
+
     case2_template = """
-            {persona}
-            The user's input contains information about a new event they want to add to their schedule. You have two tasks to perform:
+    {persona}
+    {chat_type}
+    The user's input contains information about a new event they want to add to their schedule. You have two tasks to perform:
 
-            1. Respond kindly to the user's input. YOU MUST USE {output_language} TO RESPOND TO THE USER INPUT.
-            2. Organize the event the user wants to add into a json format for saving in a database. The returned json will have keys for info, location, person, and date.
-            - info: Summarizes what the user wants to do. This value must always be present.
-            - location: If the user's event information includes a place, save that place as the value.
-            - person: If the user's event mentions a person they want to include, save that person as the value.
-            - date: If the user's event information includes a specific date and time, save that date and time in datetime format. Dates should be organized based on the current time at the user's location. Current time is {current_time}.
-            Separate the outputs for tasks 1 and 2 with a special token <separate>.
+    1. Respond kindly to the user's input. YOU MUST USE {output_language} TO RESPOND TO THE USER INPUT.
+    2. Organize the event the user wants to add into a json format for saving in a database. The returned json will have keys for info, location, person, start_time, end_time, and category.
+    - info: Summarizes what the user wants to do. This value must always be present.
+    - location: If the user's event information includes a place, save that place as the value.
+    - person: If the user's event mentions a person they want to include, save that person as the value.
+    - start_time: If the user's event information includes a specific date and time, save that date and time in datetime format. Dates should be organized based on the current time at the user's location. Current time is {current_time}.
+    - end_time: If the user's event information includes an end time, save that date and time in datetime format.
+    - category: Choose the most appropriate category for the event from the following list: {categories}.
+    Separate the outputs for tasks 1 and 2 with a special token <separate>.
 
-            Example for one-shot learning:
+    Example for one-shot learning:
 
-            User input: I have a meeting with Dr. Smith at her office on March 3rd at 10am.
+    User input: I have a meeting with Dr. Smith at her office on March 3rd from 10am to 11am.
 
-            Response to user:
-            Shall I add your meeting with Dr. Smith at her office on March 3rd at 10am to your schedule?
-            <separate>
-            {{
-            "info": "meeting with Dr. Smith",
-            "location": "Dr. Smith's office",
-            "person": "Dr. Smith",
-            "date": "2023-03-03T10:00:00"
-            }}
-            
-            User input: {question}
-            
-            Response to user:
-            """
+    Response to user:
+    Shall I add your meeting with Dr. Smith at her office on March 3rd from 10am to 11am to your schedule?
+    <separate>
+    {{
+    "info": "meeting with Dr. Smith",
+    "location": "Dr. Smith's office",
+    "person": "Dr. Smith",
+    "start_time": "2023-03-03T10:00:00",
+    "end_time": "2023-03-03T11:00:00",
+    "category": "Work"
+    }}
+
+    User input: {question}
+
+    Response to user:
+    """
 
     case3_template = """
     {persona}
+    {chat_type}
+    Current time is {current_time}. Respond to the user considering the current time.
     When responding to user inputs, it's crucial to adapt your responses to the specified output language, maintaining a consistent and accessible communication style. YOU MUST USE {output_language} TO RESPOND TO THE USER INPUT. Your responses should not only be accurate but also display empathy and understanding of the user's needs.
     You are equipped with a state-of-the-art RAG (Retrieval-Augmented Generation) technique, enabling you to dynamically pull relevant schedule information from a comprehensive database tailored to the user's specific inquiries. This technique enhances your ability to provide precise, context-aware responses by leveraging real-time data retrieval combined with advanced natural language understanding.