app.py

import gradio as gr
import psycopg2
from openai import OpenAI
import json
import os
from typing import List, Dict, Tuple, Any
from pgvector.psycopg2 import register_vector
import numpy as np
from datetime import datetime
from sklearn.preprocessing import normalize

# 가중치 및 임계값 설정
DEFAULT_FULL_WEIGHT = 0.2
DEFAULT_TOPIC_WEIGHT = 0.5
DEFAULT_CUSTOMER_WEIGHT = 0.2
DEFAULT_AGENT_WEIGHT = 0.1
DEFAULT_SIMILARITY_THRESHOLD = 0

# DB 연결 설정
def get_db_conn():
    return psycopg2.connect(
        host=os.environ["VECTOR_HOST"],
        port=5432,
        dbname=os.environ["VECTOR_DBNAME"],
        user=os.environ["VECTOR_USER"],
        password=os.environ["VECTOR_SECRET"]
    )

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 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 format_vector_for_pg(vector: List[float]) -> str:
    """벡터를 PostgreSQL 포맷으로 변환합니다."""
    return f"[{','.join(str(x) for x in vector)}]"

def search_similar_chat(query: str, max_results: int = 100) -> List[Dict]:
    """
    다중 임베딩된 채팅 데이터에서 유사한 콘텐츠를 검색합니다.
    
    Args:
        query (str): 검색할 쿼리 텍스트
        max_results (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
    
    print(f"다중 임베딩 검색 시작: 쿼리='{query}', 가중치=(full={full_w}, topic={topic_w}, customer={customer_w}, agent={agent_w}), 최대 결과={limit}")
    
    try:
        # 쿼리 임베딩 생성
        raw_embedding = np.array(get_embedding(query))
        # L2 정규화 적용
        query_embedding = normalize(raw_embedding.reshape(1, -1), norm='l2')[0]
        print(f"임베딩 정규화 전/후 첫 5개 요소: {raw_embedding[:5]} -> {query_embedding[:5]}")
        
        # Java 방식: 벡터를 문자열로 변환
        query_vector = format_vector_for_pg(query_embedding)
        
        # DB 연결
        conn = get_db_conn()
        register_vector(conn)
        
        # Java 방식: 문자열 포맷팅 사용한 SQL 쿼리
        sql = f"""
            WITH embeddings AS (
                SELECT 
                    id, 
                    metadata, 
                    content,
                    CASE WHEN full_embedding IS NOT NULL THEN (full_embedding <=> '{query_vector}'::vector) ELSE 0 END * {full_w} as full_sim,
                    CASE WHEN topic_embedding IS NOT NULL THEN (topic_embedding <=> '{query_vector}'::vector) ELSE 0 END * {topic_w} as topic_sim,
                    CASE WHEN customer_embedding IS NOT NULL THEN (customer_embedding <=> '{query_vector}'::vector) ELSE 0 END * {customer_w} as customer_sim,
                    CASE WHEN agent_embedding IS NOT NULL THEN (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 {limit}
        """
        
        with conn.cursor() as cur:
            print(f"쿼리 실행 - Java 방식 포맷팅, 가중치 설정={full_w}, {topic_w}, {customer_w}, {agent_w}, 결과 제한={limit}")
            # Java 방식: 매개변수 없이 직접 쿼리 실행
            cur.execute(sql)
            rows = cur.fetchall()
            
            print(f"검색 결과: 총 {len(rows)}개 데이터 조회됨")
            if len(rows) > 0:
                print(f"첫 번째 결과 ID: {rows[0][0]}, 유사도: {float(rows[0][3])}")
            
            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}")
                    print(f"문제가 발생한 메타데이터: {metadata_json[:200]}...")
                    continue
        
        # 임계값 필터링
        filtered_results = [r for r in results if r["similarityScore"] >= threshold]
        print(f"임계값({threshold}) 이상 결과: {len(filtered_results)}개 / 전체 {len(results)}개")
        
        if len(filtered_results) > 0:
            print(f"가장 높은 유사도 점수: {filtered_results[0]['similarityScore']}")
            print(f"상위 결과 챗ID: {filtered_results[0].get('chatId')}, 주제: {filtered_results[0].get('topic', '')[:50]}...")
        
        return filtered_results
    
    except Exception as e:
        print(f"다중 임베딩 검색 중 오류 발생: {str(e)}")
        return []
    
    finally:
        if 'conn' in locals():
            conn.close()

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): 검색할 쿼리 텍스트
        start_date (str): 검색 시작 날짜 (YYYY-MM-DD 형식)
        end_date (str): 검색 종료 날짜 (YYYY-MM-DD 형식)
        max_results (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
    
    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 []
        
        # 쿼리 임베딩 생성
        raw_embedding = np.array(get_embedding(query))
        # L2 정규화 적용
        query_embedding = normalize(raw_embedding.reshape(1, -1), norm='l2')[0]
        print(f"날짜 검색 - 임베딩 정규화 전/후 첫 5개 요소: {raw_embedding[:5]} -> {query_embedding[:5]}")
        
        # Java 방식: 벡터를 문자열로 변환
        query_vector = format_vector_for_pg(query_embedding)
        
        # DB 연결
        conn = get_db_conn()
        register_vector(conn)
        
        # Java 방식: 문자열 포맷팅 사용한 SQL 쿼리 시작
        sql = f"""
            WITH embeddings AS (
                SELECT 
                    id, 
                    metadata, 
                    content,
                    CASE WHEN full_embedding IS NOT NULL THEN (full_embedding <=> '{query_vector}'::vector) ELSE 0 END * {full_w} as full_sim,
                    CASE WHEN topic_embedding IS NOT NULL THEN (topic_embedding <=> '{query_vector}'::vector) ELSE 0 END * {topic_w} as topic_sim,
                    CASE WHEN customer_embedding IS NOT NULL THEN (customer_embedding <=> '{query_vector}'::vector) ELSE 0 END * {customer_w} as customer_sim,
                    CASE WHEN agent_embedding IS NOT NULL THEN (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
        """
        
        # 날짜 필터 추가
        if start_timestamp is not None:
            sql += f" AND (metadata->>'startTime')::bigint >= {start_timestamp}"
        
        if end_timestamp is not None:
            sql += f" AND (metadata->>'startTime')::bigint <= {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
        """
        
        with conn.cursor() as cur:
            print(f"날짜 검색 쿼리 실행: 시작일={start_date}({start_timestamp}), 종료일={end_date}({end_timestamp})")
            # 여기서는 limit만 매개변수로 전달
            cur.execute(sql, (limit,))
            rows = cur.fetchall()
            
            print(f"날짜 필터링 검색 결과: 총 {len(rows)}개 데이터 조회됨")
            if len(rows) > 0:
                print(f"첫 번째 결과 ID: {rows[0][0]}, 유사도: {float(rows[0][3])}")
            
            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}")
                    print(f"문제가 발생한 메타데이터: {metadata_json[:200]}...")
                    continue
        
        # 임계값 필터링
        filtered_results = [r for r in results if r["similarityScore"] >= threshold]
        print(f"날짜 검색 - 임계값({threshold}) 이상 결과: {len(filtered_results)}개 / 전체 {len(results)}개")
        
        if len(filtered_results) > 0:
            print(f"날짜 검색 - 가장 높은 유사도 점수: {filtered_results[0]['similarityScore']}")
            print(f"날짜 검색 - 상위 결과 챗ID: {filtered_results[0].get('chatId')}, 시작시간: {filtered_results[0].get('startTime')}")
        
        return filtered_results
    
    except Exception as e:
        print(f"다중 임베딩 날짜 검색 중 오류 발생: {str(e)}")
        return []
    
    finally:
        if 'conn' in locals():
            conn.close()

# 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")

if __name__ == "__main__":
    demo.launch(mcp_server=True)