유사도 정규화 처리

app.py

@@ -18,19 +18,17 @@ def get_db_conn():
         password=os.environ["VECTOR_SECRET"]
     )
 
-client = OpenAI()  # 환경변수에 OPENAI_API_KEY가 있으면 자동 인식
+client = OpenAI()
 
 def get_embedding(text: str) -> List[float]:
-    """
-    텍스트를 임베딩 벡터로 변환합니다.
-    """
+    """텍스트를 임베딩 벡터로 변환합니다."""
     response = client.embeddings.create(
         input=text,
         model="text-embedding-3-small"
     )
     return response.data[0].embedding
 
-def search_similar_chats(query: str, maxResults: int = 10000) -> List[Dict]:
+def search_similar_chats(query: str, maxResults: int = 200) -> List[Dict]:
     """
     유사한 채팅 문서를 검색합니다.
     Args:
@@ -39,33 +37,38 @@ def search_similar_chats(query: str, maxResults: int = 10000) -> List[Dict]:
     Returns:
         List[Dict]: 검색 결과 목록
     """
-    embedding = np.array(get_embedding(query))  # numpy array로 변환
+    embedding = np.array(get_embedding(query))
     conn = get_db_conn()
-    register_vector(conn)  # 벡터 타입 자동 변환 지원
-    with conn.cursor() as cur:
-        cur.execute("""
-            SELECT id, metadata, content, embedding <#> %s AS distance
-            FROM vector_store
-            ORDER BY embedding <#> %s
-            LIMIT %s
-        """, (embedding, embedding, maxResults))
-        rows = cur.fetchall()
-    conn.close()
-    return [
-        {
+    register_vector(conn)
+    
+    try:
+        with conn.cursor() as cur:
+            # 코사인 유사도 연산자 변경 (<=> 사용)
+            cur.execute("""
+                SELECT id, metadata, content, 
+                       1 - (embedding <=> %s) AS similarity
+                FROM vector_store
+                ORDER BY similarity DESC
+                LIMIT %s
+            """, (embedding, maxResults))
+            
+            rows = cur.fetchall()
+        return [{
             "id": row[0],
             "metadata": row[1],
             "content": row[2],
-            "distance": row[3]
-        }
-        for row in rows
-    ]
+            "similarity": float(row[3])
+        } for row in rows]
+    except Exception as e:
+        raise RuntimeError(f"DB 검색 오류: {str(e)}")
+    finally:
+        conn.close()
 
 def search_similar_chats_by_date(
     query: str,
     startDate: str = None,
     endDate: str = None,
-    maxResults: int = 10000
+    maxResults: int = 200
 ) -> List[Dict]:
     """
     지정된 날짜 범위에 해당하는 유사한 채팅 문서를 검색합니다.
@@ -78,54 +81,50 @@ def search_similar_chats_by_date(
     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")
-        except Exception as e:
-            raise ValueError(f"날짜 형식이 올바르지 않습니다.: {e}")
-    
-    embedding = np.array(get_embedding(query))  # numpy array로 변환
+    try:
+        start_dt = datetime.strptime(startDate, "%Y-%m-%d") if startDate else None
+        end_dt = datetime.strptime(endDate, "%Y-%m-%d") if endDate else None
+    except ValueError as e:
+        raise ValueError(f"날짜 형식 오류: {e}")
+
+    embedding = np.array(get_embedding(query))
     conn = get_db_conn()
-    register_vector(conn)  # 벡터 타입 자동 변환 지원
-    
-    # SQL 쿼리 구성
-    sql_query = """
-        SELECT id, metadata, content, embedding <#> %s AS distance
-        FROM vector_store
-        WHERE (metadata->>'startTime') IS NOT NULL
-          AND (metadata->>'startTime') <> ''
-    """
-    
-    params = [embedding]
-    
-    # 날짜 필터 추가
-    if startDate not in (None, ""):
-        sql_query += " AND (metadata->>'startTime')::timestamp >= %s"
-        params.append(startDate)
-    
-    if endDate not in (None, ""):
-        sql_query += " AND (metadata->>'startTime')::timestamp <= %s"
-        params.append(endDate)
-    
-    # 벡터 거리로 정렬하고 결과 제한
-    sql_query += " ORDER BY embedding <#> %s LIMIT %s"
-    params.extend([embedding, maxResults])
-    
-    with conn.cursor() as cur:
-        cur.execute(sql_query, tuple(params))
-        rows = cur.fetchall()
-    conn.close()
+    register_vector(conn)
     
-    return [
-        {
+    try:
+        with conn.cursor() as cur:
+            base_query = """
+                SELECT id, metadata, content, 
+                       1 - (embedding <=> %s) AS similarity
+                FROM vector_store
+                WHERE 1=1
+            """
+            params = [embedding]
+            
+            # 동적 쿼리 구성
+            if startDate:
+                base_query += " AND (metadata->>'startTime')::date >= %s"
+                params.append(startDate)
+            if endDate:
+                base_query += " AND (metadata->>'startTime')::date <= %s"
+                params.append(endDate)
+            
+            base_query += " ORDER BY similarity DESC LIMIT %s"
+            params.append(maxResults)
+            
+            cur.execute(base_query, tuple(params))
+            rows = cur.fetchall()
+            
+        return [{
             "id": row[0],
             "metadata": row[1],
             "content": row[2],
-            "distance": row[3]
-        }
-        for row in rows
-    ]
+            "similarity": float(row[3])
+        } for row in rows]
+    except Exception as e:
+        raise RuntimeError(f"DB 검색 오류: {str(e)}")
+    finally:
+        conn.close()
 
 # Gradio Blocks에 함수 등록
 with gr.Blocks() as demo:
@@ -134,4 +133,4 @@ with gr.Blocks() as demo:
     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)
+    demo.launch(mcp_server=True)