Update app.py

app.py

@@ -7,9 +7,13 @@ from typing import List, Dict, Tuple, Any
 from pgvector.psycopg2 import register_vector
 import numpy as np
 from datetime import datetime
-import re
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.metrics.pairwise import cosine_similarity
+
+# 가중치 및 임계값 설정
+DEFAULT_FULL_WEIGHT = 0.2
+DEFAULT_TOPIC_WEIGHT = 0.5
+DEFAULT_CUSTOMER_WEIGHT = 0.2
+DEFAULT_AGENT_WEIGHT = 0.1
+DEFAULT_SIMILARITY_THRESHOLD = 0.7
 
 # DB 연결 설정
 def get_db_conn():
@@ -31,238 +35,277 @@ def get_embedding(text: str) -> List[float]:
     )
     return response.data[0].embedding
 
-def expand_query(query: str) -> str:
-    """
-    사용자 쿼리를 확장하여 검색 품질을 개선합니다.
-    """
-    # GPT를 활용한 쿼리 확장
-    try:
-        response = client.chat.completions.create(
-            model="gpt-3.5-turbo",
-            messages=[
-                {"role": "system", "content": "당신은 검색 쿼리 확장 전문가입니다. 사용자의 쿼리를 분석하고, 이와 관련된 키워드와 질문 형태로 확장하세요."},
-                {"role": "user", "content": f"다음 검색어를 확장해주세요: '{query}'"}
-            ],
-            temperature=0.3,
-            max_tokens=150
-        )
-        expanded = query + " " + response.choices[0].message.content
-        return expanded
-    except:
-        # 오류 발생 시 원본 쿼리 반환
-        return query
+def format_vector_for_pg(vector: List[float]) -> str:
+    """벡터를 PostgreSQL 포맷으로 변환합니다."""
+    return f"[{','.join(str(x) for x in vector)}]"
 
-def extract_keywords(text: str) -> List[str]:
-    """
-    텍스트에서 중요 키워드를 추출합니다.
-    """
-    # 단순한 키워드 추출 (고급 NLP 라이브러리로 대체 가능)
-    # 불용어 제거 및 정규표현식으로 키워드 추출
-    stop_words = {'있는', '하는', '그리고', '입니다', '그것은', '있습니다', '합니다', '그런', '이런', '저런', '그냥'}
-    words = re.findall(r'\w+', text.lower())
-    keywords = [w for w in words if len(w) > 1 and w not in stop_words]
-    return list(set(keywords))
+def get_text_value(node, field_name):
+    """JSON 노드에서 텍스트 값을 안전하게 추출합니다."""
+    if node and field_name in node and node[field_name] is not None:
+        return node[field_name]
+    return None
 
-def perform_hybrid_search(
-    query: str, 
-    vector_results: List[Dict], 
-    keyword_weight: float = 0.3,
-    similarity_threshold: float = 0.4
-) -> List[Dict]:
+def search_similar_chat(query: str, max_results: int = 100) -> List[Dict]:
     """
-    벡터 검색과 키워드 검색을 결합한 하이브리드 검색을 수행합니다.
-    """
-    # 임계값 미만의 결과 필터링
-    filtered_results = [r for r in vector_results if r["similarity"] >= similarity_threshold]
-    
-    if not filtered_results:
-        # 결과가 없으면 임계값을 낮춰서 재시도
-        filtered_results = [r for r in vector_results if r["similarity"] >= similarity_threshold * 0.7]
-    
-    if not filtered_results:
-        return vector_results[:5]  # 여전히 없으면 상위 5개 반환
-    
-    # 키워드 검색 가중치 적용
-    keywords = extract_keywords(query)
+    다중 임베딩된 채팅 데이터에서 유사한 콘텐츠를 검색합니다.
     
-    for result in filtered_results:
-        content = result.get("content", "")
-        keyword_matches = sum(1 for kw in keywords if kw.lower() in content.lower())
-        keyword_score = keyword_matches / max(len(keywords), 1)
-        
-        # 최종 점수 계산 (벡터 유사도 + 키워드 가중치)
-        result["original_similarity"] = result["similarity"]
-        result["keyword_score"] = keyword_score
-        result["similarity"] = (1 - keyword_weight) * result["similarity"] + keyword_weight * keyword_score
-    
-    # 최종 점수로 재정렬
-    return sorted(filtered_results, key=lambda x: x["similarity"], reverse=True)
-
-def preprocess_query(query: str) -> str:
-    """
-    검색 쿼리를 전처리하여 검색 품질을 개선합니다.
-    """
-    # 검색에 맞게 프롬프트 재구성
-    return f"다음 질문이나 주제와 관련된 대화를 찾아주세요: {query}"
-
-def search_similar_chats(query: str, maxResults: int = 200) -> List[Dict]:
-    """
-    유사한 채팅 문서를 검색합니다.
     Args:
         query (str): 검색할 쿼리 텍스트
-        maxResults (int): 반환할 최대 결과 수
+        max_results (int): 반환할 최대 결과 수
+        
     Returns:
         List[Dict]: 검색 결과 목록
     """
-    # 쿼리 전처리 및 확장
-    processed_query = preprocess_query(query)
-    try:
-        expanded_query = expand_query(processed_query)
-    except:
-        expanded_query = processed_query
+    limit = max_results if max_results is not None else 100
+    
+    # 가중치 설정
+    full_w = DEFAULT_FULL_WEIGHT
+    topic_w = DEFAULT_TOPIC_WEIGHT
+    customer_w = DEFAULT_CUSTOMER_WEIGHT
+    agent_w = DEFAULT_AGENT_WEIGHT
+    threshold = DEFAULT_SIMILARITY_THRESHOLD
     
-    embedding = np.array(get_embedding(expanded_query))
-    conn = get_db_conn()
-    register_vector(conn)
+    print(f"다중 임베딩 검색 시작: 쿼리='{query}', 가중치=(full={full_w}, topic={topic_w}, customer={customer_w}, agent={agent_w}), 최대 결과={limit}")
     
     try:
+        # 쿼리 임베딩 생성
+        query_embedding = np.array(get_embedding(query))
+        query_vector = format_vector_for_pg(query_embedding)
+        
+        # DB 연결
+        conn = get_db_conn()
+        register_vector(conn)
+        
+        # 여러 필드를 가중치로 조합한 유사도 검색 SQL
+        sql = f"""
+            WITH embeddings AS (
+                SELECT 
+                    id, 
+                    metadata, 
+                    content,
+                    CASE WHEN full_embedding IS NOT NULL THEN 1 - (full_embedding <=> '{query_vector}'::vector) ELSE 0 END * {full_w} as full_sim,
+                    CASE WHEN topic_embedding IS NOT NULL THEN 1 - (topic_embedding <=> '{query_vector}'::vector) ELSE 0 END * {topic_w} as topic_sim,
+                    CASE WHEN customer_embedding IS NOT NULL THEN 1 - (customer_embedding <=> '{query_vector}'::vector) ELSE 0 END * {customer_w} as customer_sim,
+                    CASE WHEN agent_embedding IS NOT NULL THEN 1 - (agent_embedding <=> '{query_vector}'::vector) ELSE 0 END * {agent_w} as agent_sim
+                FROM vector_store_multi_embeddings
+                WHERE full_embedding IS NOT NULL
+                OR topic_embedding IS NOT NULL
+                OR customer_embedding IS NOT NULL
+                OR agent_embedding IS NOT NULL
+            )
+            SELECT 
+                id, 
+                metadata, 
+                content,
+                (full_sim + topic_sim + customer_sim + agent_sim) as combined_similarity
+            FROM embeddings
+            ORDER BY combined_similarity DESC
+            LIMIT %s
+        """
+        
         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))
-            
+            cur.execute(sql, (limit,))
             rows = cur.fetchall()
             
-        results = [{
-            "id": row[0],
-            "metadata": row[1],
-            "content": row[2],
-            "similarity": float(row[3])
-        } for row in rows]
+            results = []
+            for row in rows:
+                id_val = row[0]
+                metadata_json = row[1]
+                content = row[2]
+                similarity_score = float(row[3])
+                
+                # 메타데이터 파싱
+                try:
+                    metadata = json.loads(metadata_json) if isinstance(metadata_json, str) else metadata_json
+                    
+                    result = {
+                        "id": id_val,
+                        "similarityScore": similarity_score,
+                        "content": content,
+                        "chatId": get_text_value(metadata, "chatId"),
+                        "topic": get_text_value(metadata, "topic")
+                    }
+                    
+                    # 시간 필드 변환 없이 그대로 사용
+                    if "startTime" in metadata and metadata["startTime"] is not None:
+                        result["startTime"] = metadata["startTime"]
+                    
+                    if "endTime" in metadata and metadata["endTime"] is not None:
+                        result["endTime"] = metadata["endTime"]
+                    
+                    results.append(result)
+                except Exception as e:
+                    print(f"메타데이터 파싱 오류: {e}")
+                    continue
         
-        # 하이브리드 검색 적용
-        results = perform_hybrid_search(
-            query, 
-            results, 
-            keyword_weight=0.3, 
-            similarity_threshold=0.4
-        )
+        # 임계값 필터링
+        filtered_results = [r for r in results if r["similarityScore"] >= threshold]
         
-        return results
+        return filtered_results
+    
     except Exception as e:
-        raise RuntimeError(f"DB 검색 오류: {str(e)}")
+        print(f"다중 임베딩 검색 중 오류 발생: {str(e)}")
+        return []
+    
     finally:
-        conn.close()
+        if 'conn' in locals():
+            conn.close()
 
-def search_similar_chats_by_date(
-    query: str,
-    startDate: str = None,
-    endDate: str = None,
-    maxResults: int = 200
+def search_similar_chat_by_date(
+    query: str, 
+    start_date: str = None, 
+    end_date: str = None, 
+    max_results: int = 100
 ) -> List[Dict]:
     """
-    지정된 날짜 범위에 해당하는 유사한 채팅 문서를 검색합니다.
+    지정된 날짜 범위 내의 다중 임베딩 채팅 데이터를 검색합니다.
     
     Args:
-        query (str): 검색 쿼리
-        startDate (str): 검색 시작 날짜 (YYYY-MM-DD)
-        endDate (str): 검색 종료 날짜 (YYYY-MM-DD)
-        maxResults (int): 반환할 최대 결과 수
+        query (str): 검색할 쿼리 텍스트
+        start_date (str): 검색 시작 날짜 (YYYY-MM-DD 형식)
+        end_date (str): 검색 종료 날짜 (YYYY-MM-DD 형식)
+        max_results (int): 반환할 최대 결과 수
+        
     Returns:
         List[Dict]: 검색 결과 목록
     """
-    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}")
-
-    # 쿼리 전처리 및 확장
-    processed_query = preprocess_query(query)
-    try:
-        expanded_query = expand_query(processed_query)
-    except:
-        expanded_query = processed_query
-        
-    embedding = np.array(get_embedding(expanded_query))
-    conn = get_db_conn()
-    register_vector(conn)
+    limit = max_results if max_results is not None else 100
+    
+    # 가중치 설정
+    full_w = DEFAULT_FULL_WEIGHT
+    topic_w = DEFAULT_TOPIC_WEIGHT
+    customer_w = DEFAULT_CUSTOMER_WEIGHT
+    agent_w = DEFAULT_AGENT_WEIGHT
+    threshold = DEFAULT_SIMILARITY_THRESHOLD
+    
+    print(f"다중 임베딩 날짜 검색 시작: 쿼리='{query}', 시작일={start_date}, 종료일={end_date}, 최대 결과={limit}")
     
     try:
+        # 날짜 필터 파라미터 생성
+        start_timestamp = None
+        end_timestamp = None
+        
+        if start_date and start_date.strip():
+            try:
+                start_datetime = datetime.strptime(start_date, '%Y-%m-%d')
+                start_timestamp = int(start_datetime.timestamp() * 1000)  # 밀리초 단위로 변환
+            except ValueError as e:
+                print(f"시작 날짜 형식 오류: {str(e)}")
+                return []
+        
+        if end_date and end_date.strip():
+            try:
+                # 종료일의 23:59:59로 설정
+                end_datetime = datetime.strptime(end_date + ' 23:59:59', '%Y-%m-%d %H:%M:%S')
+                end_timestamp = int(end_datetime.timestamp() * 1000)  # 밀리초 단위로 변환
+            except ValueError as e:
+                print(f"종료 날짜 형식 오류: {str(e)}")
+                return []
+        
+        # 쿼리 임베딩 생성
+        query_embedding = np.array(get_embedding(query))
+        query_vector = format_vector_for_pg(query_embedding)
+        
+        # DB 연결
+        conn = get_db_conn()
+        register_vector(conn)
+        
+        # 여러 필드를 가중치로 조합한 유사도 검색 SQL (날짜 필터 추가)
+        sql = f"""
+            WITH embeddings AS (
+                SELECT 
+                    id, 
+                    metadata, 
+                    content,
+                    CASE WHEN full_embedding IS NOT NULL THEN 1 - (full_embedding <=> '{query_vector}'::vector) ELSE 0 END * {full_w} as full_sim,
+                    CASE WHEN topic_embedding IS NOT NULL THEN 1 - (topic_embedding <=> '{query_vector}'::vector) ELSE 0 END * {topic_w} as topic_sim,
+                    CASE WHEN customer_embedding IS NOT NULL THEN 1 - (customer_embedding <=> '{query_vector}'::vector) ELSE 0 END * {customer_w} as customer_sim,
+                    CASE WHEN agent_embedding IS NOT NULL THEN 1 - (agent_embedding <=> '{query_vector}'::vector) ELSE 0 END * {agent_w} as agent_sim
+                FROM vector_store_multi_embeddings
+                WHERE full_embedding IS NOT NULL
+                OR topic_embedding IS NOT NULL
+                OR customer_embedding IS NOT NULL
+                OR agent_embedding IS NOT NULL
+        """
+        
+        params = []
+        
+        # 날짜 필터 추가
+        if start_timestamp is not None:
+            sql += f" AND (metadata->>'startTime')::bigint >= %s"
+            params.append(start_timestamp)
+        
+        if end_timestamp is not None:
+            sql += f" AND (metadata->>'startTime')::bigint <= %s"
+            params.append(end_timestamp)
+        
+        sql += """
+            )
+            SELECT 
+                id, 
+                metadata, 
+                content,
+                (full_sim + topic_sim + customer_sim + agent_sim) as combined_similarity
+            FROM embeddings
+            ORDER BY combined_similarity DESC
+            LIMIT %s
+        """
+        
+        params.append(limit)
+        
         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))
+            cur.execute(sql, tuple(params))
             rows = cur.fetchall()
             
-        results = [{
-            "id": row[0],
-            "metadata": row[1],
-            "content": row[2],
-            "similarity": float(row[3])
-        } for row in rows]
-        
-        # 하이브리드 검색 적용
-        results = perform_hybrid_search(
-            query, 
-            results, 
-            keyword_weight=0.3, 
-            similarity_threshold=0.4
-        )
-            
-        # 메타데이터 기반 가중치 적용
-        keywords = extract_keywords(query)
-        for result in results:
-            metadata = result.get("metadata", {})
-            if not metadata or isinstance(metadata, str):
-                continue
+            results = []
+            for row in rows:
+                id_val = row[0]
+                metadata_json = row[1]
+                content = row[2]
+                similarity_score = float(row[3])
                 
-            # 주제(topic) 필드에 키워드가 있는지 확인
-            topic = metadata.get("topic", "")
-            topic_matches = sum(1 for kw in keywords if kw.lower() in topic.lower())
-            
-            # 주제 일치 가중치 적용
-            if topic_matches > 0:
-                topic_boost = 0.1 * min(topic_matches, 3)  # 최대 0.3 가중치
-                result["similarity"] += topic_boost
-                result["topic_boost"] = topic_boost
+                # 메타데이터 파싱
+                try:
+                    metadata = json.loads(metadata_json) if isinstance(metadata_json, str) else metadata_json
+                    
+                    result = {
+                        "id": id_val,
+                        "similarityScore": similarity_score,
+                        "content": content,
+                        "chatId": get_text_value(metadata, "chatId"),
+                        "topic": get_text_value(metadata, "topic")
+                    }
+                    
+                    # 시간 필드 변환 없이 그대로 사용
+                    if "startTime" in metadata and metadata["startTime"] is not None:
+                        result["startTime"] = metadata["startTime"]
+                    
+                    if "endTime" in metadata and metadata["endTime"] is not None:
+                        result["endTime"] = metadata["endTime"]
+                    
+                    results.append(result)
+                except Exception as e:
+                    print(f"메타데이터 파싱 오류: {e}")
+                    continue
         
-        # 결과 재정렬
-        results = sorted(results, key=lambda x: x["similarity"], reverse=True)
+        # 임계값 필터링
+        filtered_results = [r for r in results if r["similarityScore"] >= threshold]
         
-        return results
+        return filtered_results
+    
     except Exception as e:
-        raise RuntimeError(f"DB 검색 오류: {str(e)}")
+        print(f"다중 임베딩 날짜 검색 중 오류 발생: {str(e)}")
+        return []
+    
     finally:
-        conn.close()
+        if 'conn' in locals():
+            conn.close()
 
 # Gradio Blocks에 함수 등록
 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_similar_chats_by_date, inputs=["text", "text", "text", "number"], outputs="json", api_name="search_similar_chats_by_date")
+    gr.Interface(fn=search_similar_chat, inputs=["text", "number"], outputs="json", api_name="search_similar_chat")
+    gr.Interface(fn=search_similar_chat_by_date, inputs=["text", "text", "text", "number"], outputs="json", api_name="search_similar_chat_by_date")
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(mcp_server=True)