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

# 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()  # 환경변수에 OPENAI_API_KEY가 있으면 자동 인식

def get_embedding(text: str) -> List[float]:
    """
    텍스트를 임베딩 벡터로 변환합니다.
    """
    response = client.embeddings.create(
        input=text,
        model="text-embedding-ada-002"
    )
    return response.data[0].embedding

def search_similar_chats(query: str, maxResults: int = 10) -> List[Dict]:
    """
    유사한 채팅 문서를 검색합니다.

    Args:
        query (str): 검색할 쿼리 텍스트
        maxResults (int): 반환할 최대 결과 수

    Returns:
        List[Dict]: 검색 결과 목록
    """
    embedding = np.array(get_embedding(query))  # numpy array로 변환
    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
            WHERE metadata->>'documentType' = 'chatAnalysis'
            ORDER BY embedding <#> %s
            LIMIT %s
        """, (embedding, embedding, maxResults))
        rows = cur.fetchall()
    conn.close()
    return [
        {
            "id": row[0],
            "metadata": row[1],
            "content": row[2],
            "distance": row[3]
        }
        for row in rows
    ]

def search_chats_by_category(category: str, maxResults: int = 10) -> List[Dict]:
    """
    특정 카테고리의 채팅 문서를 검색합니다.

    Args:
        category (str): 카테고리명
        maxResults (int): 반환할 최대 결과 수

    Returns:
        List[Dict]: 검색 결과 목록
    """
    conn = get_db_conn()
    with conn.cursor() as cur:
        cur.execute("""
            SELECT id, metadata, content
            FROM vector_store
            WHERE metadata->>'documentType' = 'chatAnalysis'
              AND metadata->>'category' = %s
            LIMIT %s
        """, (category, maxResults))
        rows = cur.fetchall()
    conn.close()
    return [
        {
            "id": row[0],
            "metadata": row[1],
            "content": row[2]
        }
        for row in rows
    ]

def search_chats_by_date(startDate: str = None, endDate: str = None, maxResults: int = 10) -> List[Dict]:
    """
    지정된 날짜 범위 내의 채팅 문서를 검색합니다.

    Args:
        startDate (str): 검색 시작 날짜 (YYYY-MM-DD)
        endDate (str): 검색 종료 날짜 (YYYY-MM-DD)
        maxResults (int): 반환할 최대 결과 수

    Returns:
        List[Dict]: 검색 결과 목록
    """
    conn = get_db_conn()
    query = """
        SELECT id, metadata, content
        FROM vector_store
        WHERE metadata->>'documentType' = 'chatAnalysis'
    """
    params = []
    if startDate not in (None, ""):
        query += " AND (metadata->>'startTime')::timestamp >= %s"
        params.append(startDate)
    if endDate not in (None, ""):
        query += " AND (metadata->>'startTime')::timestamp < %s"
        params.append(endDate)
    query += " LIMIT %s"
    params.append(maxResults)
    with conn.cursor() as cur:
        cur.execute(query, tuple(params))
        rows = cur.fetchall()
    conn.close()
    return [
        {
            "id": row[0],
            "metadata": row[1],
            "content": row[2]
        }
        for row in rows
    ]

# Gradio Blocks에 함수 등록
with gr.Blocks() as demo:
    gr.Markdown("# MCP ToolService 예시")
    gr.Interface(fn=search_similar_chats, inputs=["text", "number"], outputs="json", title="search_similar_chats")
    gr.Interface(fn=search_chats_by_category, inputs=["text", "number"], outputs="json", title="search_chats_by_category")
    gr.Interface(fn=search_chats_by_date, inputs=["text", "text", "number"], outputs="json", title="search_chats_by_date")

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