test

app.py

@@ -3,16 +3,13 @@ import psycopg2
 from openai import OpenAI
 import json
 import os
-from typing import List, Dict
+from typing import List, Dict, Tuple, Any
 from pgvector.psycopg2 import register_vector
 import numpy as np
-
-# 가중치 및 임계값 설정
-DEFAULT_FULL_WEIGHT = 0.2
-DEFAULT_TOPIC_WEIGHT = 0.5
-DEFAULT_CUSTOMER_WEIGHT = 0.2
-DEFAULT_AGENT_WEIGHT = 0.1
-DEFAULT_SIMILARITY_THRESHOLD = 0.5
+from datetime import datetime
+import re
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
 
 # DB 연결 설정
 def get_db_conn():
@@ -24,294 +21,248 @@ def get_db_conn():
         password=os.environ["VECTOR_SECRET"]
     )
 
-# OpenAI 클라이언트 초기화
 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 expand_query(query: str) -> str:
     """
-    텍스트를 OpenAI의 text-embedding-ada-002 모델을 사용하여 임베딩 벡터로 변환합니다.
-    Java의 float[](float32)와 호환되도록 명시적으로 float32로 변환합니다.
-    
-    Args:
-        text (str): 임베딩할 텍스트
-    
-    Returns:
-        List[float]: 임베딩 벡터 (float32)
+    사용자 쿼리를 확장하여 검색 품질을 개선합니다.
     """
+    # GPT를 활용한 쿼리 확장
     try:
-        response = client.embeddings.create(
-            input=text,
-            model="text-embedding-ada-002",
-            encoding_format="float"
+        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
         )
-        # 명시적으로 float32로 변환하여 Java의 float[]와 호환되게 함
-        return np.array(response.data[0].embedding, dtype=np.float32).tolist()
-    except Exception as e:
-        print(f"임베딩 생성 중 오류 발생: {str(e)}")
-        raise
+        expanded = query + " " + response.choices[0].message.content
+        return expanded
+    except:
+        # 오류 발생 시 원본 쿼리 반환
+        return query
 
-def format_vector_for_pg(vector: List[float]) -> str:
+def extract_keywords(text: str) -> List[str]:
     """
-    임베딩 벡터를 PostgreSQL 포맷으로 변환합니다.
-    입력된 벡터가 float32 타입인지 확인합니다.
+    텍스트에서 중요 키워드를 추출합니다.
     """
-    # 벡터가 float32 타입인지 확인하고, 아니면 변환
-    # NumPy 배열이 아닌 경우에도 처리
-    if not isinstance(vector, np.ndarray):
-        vector = np.array(vector, dtype=np.float32)
-    elif vector.dtype != np.float32:
-        vector = vector.astype(np.float32)
-    vector_str = ','.join([f"{x}" for x in vector])
-    return f"[{vector_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 perform_hybrid_search(
+    query: str, 
+    vector_results: List[Dict], 
+    keyword_weight: float = 0.3,
+    similarity_threshold: float = 0.4
+) -> 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 get_text_value(node: Dict, field_name: str) -> str:
+def preprocess_query(query: str) -> str:
     """
-    딕셔너리에서 텍스트 값을 안전하게 추출합니다.
-    자바의 getTextValue() 메소드와 동일한 기능입니다.
+    검색 쿼리를 전처리하여 검색 품질을 개선합니다.
     """
-    if node and field_name in node and node[field_name] is not None:
-        return node[field_name]
-    return None
+    # 검색에 맞게 프롬프트 재구성
+    return f"다음 질문이나 주제와 관련된 대화를 찾아주세요: {query}"
 
-def search_similar_chat(query: str, max_results: int = 100) -> List[Dict]:
+def search_similar_chats(query: str, maxResults: int = 200) -> List[Dict]:
     """
-    채팅 데이터에서 유사한 콘텐츠를 검색합니다.
-    
+    유사한 채팅 문서를 검색합니다.
     Args:
         query (str): 검색할 쿼리 텍스트
-        max_results (int): 반환할 최대 결과 수
-        
+        maxResults (int): 반환할 최대 결과 수
     Returns:
         List[Dict]: 검색 결과 목록
     """
-    limit = max_results if max_results is not None else 100
+    # 쿼리 전처리 및 확장
+    processed_query = preprocess_query(query)
+    try:
+        expanded_query = expand_query(processed_query)
+    except:
+        expanded_query = processed_query
     
-    # 자바와 동일한 가중치 설정
-    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)
     
     try:
-        # 쿼리 임베딩 생성
-        query_embedding = get_embedding(query)
-        
-        # PostgreSQL 포맷으로 벡터 변환
-        query_vector = format_vector_for_pg(query_embedding)
-        
-        # DB 연결
-        conn = get_db_conn()
-        register_vector(conn)
-        
-        # 자바 코드와 동일한 SQL 쿼리 구현
-        sql = """
-            WITH embeddings AS (
-                SELECT 
-                    id, 
-                    metadata, 
-                    content,
-                    CASE WHEN full_embedding IS NOT NULL THEN 1 - (full_embedding <=> '%s'::vector) ELSE 0 END * %f as full_sim,
-                    CASE WHEN topic_embedding IS NOT NULL THEN 1 - (topic_embedding <=> '%s'::vector) ELSE 0 END * %f as topic_sim,
-                    CASE WHEN customer_embedding IS NOT NULL THEN 1 - (customer_embedding <=> '%s'::vector) ELSE 0 END * %f as customer_sim,
-                    CASE WHEN agent_embedding IS NOT NULL THEN 1 - (agent_embedding <=> '%s'::vector) ELSE 0 END * %f 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
-        """ % (query_vector, full_w, query_vector, topic_w, query_vector, customer_w, query_vector, agent_w, limit)
-        
         with conn.cursor() as cur:
-            cur.execute(sql)
-            rows = cur.fetchall()
+            # 코사인 유사도 계산
+            cur.execute("""
+                SELECT id, metadata, content, 
+                       1 - (embedding <=> %s) AS similarity
+                FROM vector_store
+                ORDER BY similarity DESC
+                LIMIT %s
+            """, (embedding, maxResults))
             
-            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
-            
-            # 임계값 필터링
-            filtered_results = [r for r in results if r["similarityScore"] >= threshold]
+            rows = cur.fetchall()
             
-            return filtered_results
-    
+        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
+        )
+        
+        return results
     except Exception as e:
-        print(f"다중 임베딩 검색 중 오류 발생: {str(e)}")
-        return []
-    
+        raise RuntimeError(f"DB 검색 오류: {str(e)}")
     finally:
-        if 'conn' in locals():
-            conn.close()
+        conn.close()
 
-def search_similar_chat_by_date(
-    query: str, 
-    start_date: str = None, 
-    end_date: str = None, 
-    max_results: int = 100
+def search_similar_chats_by_date(
+    query: str,
+    startDate: str = None,
+    endDate: str = None,
+    maxResults: int = 200
 ) -> List[Dict]:
     """
-    지정된 날짜 범위 내의 채팅 데이터를 검색합니다.
+    지정된 날짜 범위에 해당하는 유사한 채팅 문서를 검색합니다.
     
     Args:
-        query (str): 검색할 쿼리 텍스트
-        start_date (str): 검색 시작 날짜 (YYYY-MM-DD 형식)
-        end_date (str): 검색 종료 날짜 (YYYY-MM-DD 형식)
-        max_results (int): 반환할 최대 결과 수
-        
+        query (str): 검색 쿼리
+        startDate (str): 검색 시작 날짜 (YYYY-MM-DD)
+        endDate (str): 검색 종료 날짜 (YYYY-MM-DD)
+        maxResults (int): 반환할 최대 결과 수
     Returns:
         List[Dict]: 검색 결과 목록
     """
-    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
-    
     try:
-        # 쿼리 임베딩 생성
-        query_embedding = get_embedding(query)
-        
-        # PostgreSQL 포맷으로 벡터 변환
-        query_vector = format_vector_for_pg(query_embedding)
-        
-        # DB 연결
-        conn = get_db_conn()
-        register_vector(conn)
-        
-        # 자바 코드와 동일한 SQL 쿼리 시작
-        sql = """
-            WITH embeddings AS (
-                SELECT 
-                    id, 
-                    metadata, 
-                    content,
-                    CASE WHEN full_embedding IS NOT NULL THEN 1 - (full_embedding <=> '%s'::vector) ELSE 0 END * %f as full_sim,
-                    CASE WHEN topic_embedding IS NOT NULL THEN 1 - (topic_embedding <=> '%s'::vector) ELSE 0 END * %f as topic_sim,
-                    CASE WHEN customer_embedding IS NOT NULL THEN 1 - (customer_embedding <=> '%s'::vector) ELSE 0 END * %f as customer_sim,
-                    CASE WHEN agent_embedding IS NOT NULL THEN 1 - (agent_embedding <=> '%s'::vector) ELSE 0 END * %f 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
-        """ % (query_vector, full_w, query_vector, topic_w, query_vector, customer_w, query_vector, agent_w)
-        
-        # 날짜 필터 추가
-        if start_date and start_date.strip():
-            # 시작 시간 추가하여 ISO 형식으로 비교
-            iso_start_date = start_date + "T00:00:00"
-            sql += f" AND metadata->>'startTime' >= '{iso_start_date}'"
-        
-        if end_date and end_date.strip():
-            # 종료 시간 추가하여 ISO 형식으로 비교
-            iso_end_date = end_date + "T23:59:59"
-            sql += f" AND metadata->>'startTime' <= '{iso_end_date}'"
-        
-        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
-        """
+        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)
+    
+    try:
         with conn.cursor() as cur:
-            # 여기서는 limit를 파라미터로 전달
-            cur.execute(sql, (limit,))
+            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()
             
-            results = []
-            for row in rows:
-                id_val = row[0]
-                metadata_json = row[1]
-                content = row[2]
-                similarity_score = float(row[3])
+        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
                 
-                # 메타데이터 파싱
-                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")
-                    }
-                    
-                    # 시간 필드 변환 없이 그대로 사용 (이미 KST로 저장되어 있음)
-                    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
+            # 주제(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
         
-        # 임계값 필터링 (자바 코드와 동일하게 구현)
-        filtered_results = [r for r in results if r["similarityScore"] >= threshold]
+        # 결과 재정렬
+        results = sorted(results, key=lambda x: x["similarity"], reverse=True)
         
-        return filtered_results
-    
+        return results
     except Exception as e:
-        print(f"다중 임베딩 날짜 검색 중 오류 발생: {str(e)}")
-        return []
-    
+        raise RuntimeError(f"DB 검색 오류: {str(e)}")
     finally:
-        if 'conn' in locals():
-            conn.close()
+        conn.close()
 
-# Gradio 웹 인터페이스 설정
+# Gradio Blocks에 함수 등록
 with gr.Blocks() as demo:
     gr.Markdown("# Chat Analysis Search")
-    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")
+    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")
 
 if __name__ == "__main__":
     demo.launch(mcp_server=True)