app.py

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

# 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 search_similar_chats(query: str, maxResults: int = 200) -> List[Dict]:
    """
    유사한 채팅 문서를 검색합니다.
    Args:
        query (str): 검색할 쿼리 텍스트
        maxResults (int): 반환할 최대 결과 수
    Returns:
        List[Dict]: 검색 결과 목록
    """
    embedding = np.array(get_embedding(query))
    conn = get_db_conn()
    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],
            "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 = 200
) -> List[Dict]:
    """
    지정된 날짜 범위에 해당하는 유사한 채팅 문서를 검색합니다.
    
    Args:
        query (str): 검색 쿼리
        startDate (str): 검색 시작 날짜 (YYYY-MM-DD)
        endDate (str): 검색 종료 날짜 (YYYY-MM-DD)
        maxResults (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}")

    embedding = np.array(get_embedding(query))
    conn = get_db_conn()
    register_vector(conn)
    
    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],
            "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:
    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")

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