Update app.py

app.py

@@ -26,18 +26,16 @@ def get_embedding(text: str) -> List[float]:
     """
     response = client.embeddings.create(
         input=text,
-        model="text-embedding-ada-002"
+        model="text-embedding-3-small"
     )
     return response.data[0].embedding
 
-def search_similar_chats(query: str, maxResults: int = 30000) -> List[Dict]:
+def search_similar_chats(query: str, maxResults: int = 10000) -> List[Dict]:
     """
     유사한 채팅 문서를 검색합니다.
-
     Args:
         query (str): 검색할 쿼리 텍스트
         maxResults (int): 반환할 최대 결과 수
-
     Returns:
         List[Dict]: 검색 결과 목록
     """
@@ -48,7 +46,6 @@ def search_similar_chats(query: str, maxResults: int = 30000) -> List[Dict]:
         cur.execute("""
             SELECT id, metadata, content, embedding <#> %s AS distance
             FROM vector_store
-            WHERE metadata->>'documentType' = 'chatAnalysis'
             ORDER BY embedding <#> %s
             LIMIT %s
         """, (embedding, embedding, maxResults))
@@ -64,22 +61,20 @@ def search_similar_chats(query: str, maxResults: int = 30000) -> List[Dict]:
         for row in rows
     ]
 
-def search_chats_by_category_and_date(
-    category: str,
+def search_similar_chats_by_date(
+    query: str,
     startDate: str = None,
     endDate: str = None,
-    maxResults: int = 30000
+    maxResults: int = 10000
 ) -> List[Dict]:
     """
-    카테고리와 지정된 날짜 범위에 해당하는 채팅 문서를 검색합니다.
-    (검색 기간은 최대 31일까지 허용)
-
+    지정된 날짜 범위에 해당하는 유사한 채팅 문서를 검색합니다.
+    
     Args:
-        category (str): 카테고리명
+        query (str): 검색 쿼리
         startDate (str): 검색 시작 날짜 (YYYY-MM-DD)
         endDate (str): 검색 종료 날짜 (YYYY-MM-DD)
         maxResults (int): 반환할 최대 결과 수
-
     Returns:
         List[Dict]: 검색 결과 목록
     """
@@ -87,97 +82,47 @@ def search_chats_by_category_and_date(
         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)
+            raise ValueError(f"날짜 형식이 올바르지 않습니다.: {e}")
+    
+    embedding = np.array(get_embedding(query))  # numpy array로 변환
     conn = get_db_conn()
-    with conn.cursor() as cur:
-        cur.execute(query, tuple(params))
-        rows = cur.fetchall()
-    conn.close()
-    return [
-        {
-            "id": row[0],
-            "metadata": row[1],
-            "content": row[2]
-        }
-        for row in rows
-    ]
-
-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): 반환할 최대 결과 수
-
-    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
+    register_vector(conn)  # 벡터 타입 자동 변환 지원
+    
+    # SQL 쿼리 구성
+    sql_query = """
+        SELECT id, metadata, content, embedding <#> %s AS distance
         FROM vector_store
-        WHERE metadata->>'documentType' = 'chatAnalysis'
-          AND (metadata->>'startTime') IS NOT NULL
+        WHERE (metadata->>'startTime') IS NOT NULL
           AND (metadata->>'startTime') <> ''
     """
-    params = []
-    if productName not in (None, ""):
-        query += " AND (metadata->>'productName') = %s"
-        params.append(productName)
+    
+    params = [embedding]
+    
+    # 날짜 필터 추가
     if startDate not in (None, ""):
-        query += " AND (metadata->>'startTime')::timestamp >= %s"
+        sql_query += " AND (metadata->>'startTime')::timestamp >= %s"
         params.append(startDate)
+    
     if endDate not in (None, ""):
-        query += " AND (metadata->>'startTime')::timestamp <= %s"
+        sql_query += " AND (metadata->>'startTime')::timestamp <= %s"
         params.append(endDate)
-    query += " LIMIT %s"
-    params.append(maxResults)
-    conn = get_db_conn()
+    
+    # 벡터 거리로 정렬하고 결과 제한
+    sql_query += " ORDER BY embedding <#> %s LIMIT %s"
+    params.extend([embedding, maxResults])
+    
     with conn.cursor() as cur:
-        cur.execute(query, tuple(params))
+        cur.execute(sql_query, tuple(params))
         rows = cur.fetchall()
     conn.close()
+    
     return [
         {
             "id": row[0],
             "metadata": row[1],
-            "content": row[2]
+            "content": row[2],
+            "distance": row[3]
         }
         for row in rows
     ]
@@ -186,8 +131,7 @@ def search_chats_by_product_and_date(
 with gr.Blocks() as demo:
     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_and_date, inputs=["text", "text", "text", "number"], outputs="json", api_name="search_chats_by_category_and_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_similar_chats_by_date, inputs=["text", "text", "text", "number"], outputs="json", api_name="search_similar_chats_by_date")
 
 if __name__ == "__main__":
     demo.launch(mcp_server=True)