Update app.py

app.py

@@ -6,6 +6,7 @@ import os
 from typing import List, Dict
 from pgvector.psycopg2 import register_vector
 import numpy as np
+from datetime import datetime
 
 # DB 연결 설정
 def get_db_conn():
@@ -29,7 +30,7 @@ def get_embedding(text: str) -> List[float]:
     )
     return response.data[0].embedding
 
-def search_similar_chats(query: str, maxResults: int = 10) -> List[Dict]:
+def search_similar_chats(query: str, maxResults: int = 30000) -> List[Dict]:
     """
     유사한 채팅 문서를 검색합니다.
 
@@ -63,26 +64,53 @@ def search_similar_chats(query: str, maxResults: int = 10) -> List[Dict]:
         for row in rows
     ]
 
-def search_chats_by_category(category: str, maxResults: int = 10) -> List[Dict]:
+def search_chats_by_category_and_date(
+    category: str,
+    startDate: str = None,
+    endDate: str = None,
+    maxResults: int = 30000
+) -> List[Dict]:
     """
-    특정 카테고리의 채팅 문서를 검색합니다.
+    카테고리와 지정된 날짜 범위에 해당하는 채팅 문서를 검색합니다.
+    (검색 기간은 최대 31일까지 허용)
 
     Args:
         category (str): 카테고리명
+        startDate (str): 검색 시작 날짜 (YYYY-MM-DD)
+        endDate (str): 검색 종료 날짜 (YYYY-MM-DD)
         maxResults (int): 반환할 최대 결과 수
 
     Returns:
         List[Dict]: 검색 결과 목록
     """
+    if startDate not in (None, "") and endDate not in (None, ""):
+        try:
+            start_dt = datetime.strptime(startDate, "%Y-%m-%d")
+            end_dt = datetime.strptime(endDate, "%Y-%m-%d")
+            if (end_dt - start_dt).days > 31:
+                raise ValueError("검색 기간은 최대 31일까지 허용됩니다.")
+        except Exception as e:
+            raise ValueError(f"날짜 형식이 올바르지 않거나, 기간이 너무 깁니다: {e}")
+    query = """
+        SELECT id, metadata, content
+        FROM vector_store
+        WHERE metadata->>'documentType' = 'chatAnalysis'
+          AND metadata->>'category' = %s
+          AND (metadata->>'startTime') IS NOT NULL
+          AND (metadata->>'startTime') <> ''
+    """
+    params = [category]
+    if startDate not in (None, ""):
+        query += " AND (metadata->>'startTime')::timestamp >= %s"
+        params.append(startDate)
+    if endDate not in (None, ""):
+        query += " AND (metadata->>'startTime')::timestamp < %s"
+        params.append(endDate)
+    query += " LIMIT %s"
+    params.append(maxResults)
     conn = get_db_conn()
     with conn.cursor() as cur:
-        cur.execute("""
-            SELECT id, metadata, content
-            FROM vector_store
-            WHERE metadata->>'documentType' = 'chatAnalysis'
-              AND metadata->>'category' = %s
-            LIMIT %s
-        """, (category, maxResults))
+        cur.execute(query, tuple(params))
         rows = cur.fetchall()
     conn.close()
     return [
@@ -94,11 +122,18 @@ def search_chats_by_category(category: str, maxResults: int = 10) -> List[Dict]:
         for row in rows
     ]
 
-def search_chats_by_date(startDate: str = None, endDate: str = None, maxResults: int = 10) -> List[Dict]:
+def search_chats_by_product_and_date(
+    productName: str = None,
+    startDate: str = None,
+    endDate: str = None,
+    maxResults: int = 30000
+) -> List[Dict]:
     """
-    지정된 날짜 범위 내의 채팅 문서를 검색합니다.
+    상품명과 지정된 날짜 범위에 해당하는 채팅 문서를 검색합니다.
+    (검색 기간은 최대 31일까지 허용)
 
     Args:
+        productName (str): 상품명(옵션)
         startDate (str): 검색 시작 날짜 (YYYY-MM-DD)
         endDate (str): 검색 종료 날짜 (YYYY-MM-DD)
         maxResults (int): 반환할 최대 결과 수
@@ -106,7 +141,14 @@ def search_chats_by_date(startDate: str = None, endDate: str = None, maxResults:
     Returns:
         List[Dict]: 검색 결과 목록
     """
-    conn = get_db_conn()
+    if startDate not in (None, "") and endDate not in (None, ""):
+        try:
+            start_dt = datetime.strptime(startDate, "%Y-%m-%d")
+            end_dt = datetime.strptime(endDate, "%Y-%m-%d")
+            if (end_dt - start_dt).days > 31:
+                raise ValueError("검색 기간은 최대 31일까지 허용됩니다.")
+        except Exception as e:
+            raise ValueError(f"날짜 형식이 올바르지 않거나, 기간이 너무 깁니다: {e}")
     query = """
         SELECT id, metadata, content
         FROM vector_store
@@ -115,6 +157,9 @@ def search_chats_by_date(startDate: str = None, endDate: str = None, maxResults:
           AND (metadata->>'startTime') <> ''
     """
     params = []
+    if productName not in (None, ""):
+        query += " AND (metadata->>'productName') = %s"
+        params.append(productName)
     if startDate not in (None, ""):
         query += " AND (metadata->>'startTime')::timestamp >= %s"
         params.append(startDate)
@@ -123,6 +168,7 @@ def search_chats_by_date(startDate: str = None, endDate: str = None, maxResults:
         params.append(endDate)
     query += " LIMIT %s"
     params.append(maxResults)
+    conn = get_db_conn()
     with conn.cursor() as cur:
         cur.execute(query, tuple(params))
         rows = cur.fetchall()
@@ -138,10 +184,10 @@ def search_chats_by_date(startDate: str = None, endDate: str = None, maxResults:
 
 # Gradio Blocks에 함수 등록
 with gr.Blocks() as demo:
-    gr.Markdown("# MCP ToolService 예시")
+    gr.Markdown("# Chat Analysis Search")
     gr.Interface(fn=search_similar_chats, inputs=["text", "number"], outputs="json", api_name="search_similar_chats")
-    gr.Interface(fn=search_chats_by_category, inputs=["text", "number"], outputs="json", api_name="search_chats_by_category")
-    gr.Interface(fn=search_chats_by_date, inputs=["text", "text", "number"], outputs="json", api_name="search_chats_by_date")
+    gr.Interface(fn=search_chats_by_product_and_date, inputs=["text", "text", "text", "number"], outputs="json", api_name="search_chats_by_product_and_date")
+    gr.Interface(fn=search_chats_by_category_and_date, inputs=["text", "text", "text", "number"], outputs="json", api_name="search_chats_by_category_and_date")
 
 if __name__ == "__main__":
     demo.launch(mcp_server=True)