Update app.py

app.py

@@ -1,28 +1,141 @@
 import gradio as gr
+import psycopg2
+import openai
+import json
+import os
+from typing import List, Dict
 
-def letter_counter(word, letter):
+# 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"]
+    )
+
+# 임베딩 함수 (OpenAI API 예시)
+def get_embedding(text: str) -> List[float]:
+    """
+    텍스트를 임베딩 벡터로 변환합니다.
     """
-    Count the number of occurrences of a letter in a word or text.
+    response = openai.Embedding.create(
+        input=text,
+        model="text-embedding-3-large"
+    )
+    return response['data'][0]['embedding']
+
+def search_similar_chats(query: str, maxResults: int = 10) -> List[Dict]:
+    """
+    유사한 채팅 문서를 검색합니다.
 
     Args:
-        word (str): The input text to search through
-        letter (str): The letter to search for
+        query (str): 검색할 쿼리 텍스트
+        maxResults (int): 반환할 최대 결과 수
 
     Returns:
-        str: A message indicating how many times the letter appears
+        List[Dict]: 검색 결과 목록
+    """
+    embedding = get_embedding(query)
+    conn = get_db_conn()
+    with conn.cursor() as cur:
+        cur.execute("""
+            SELECT id, metadata, content, embedding <#> %s AS distance
+            FROM chat_vector_table
+            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 chat_vector_table
+            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 chat_vector_table
+        WHERE metadata->>'documentType' = 'chatAnalysis'
     """
-    word = word.lower()
-    letter = letter.lower()
-    count = word.count(letter)
-    return count
+    params = []
+    if startDate:
+        query += " AND (metadata->>'startTime')::timestamp >= %s"
+        params.append(startDate)
+    if endDate:
+        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
+    ]
 
-demo = gr.Interface(
-    fn=letter_counter,
-    inputs=["textbox", "textbox"],
-    outputs="number",
-    title="Letter Counter",
-    description="Enter text and a letter to count how many times the letter appears in the text."
-)
+# Gradio Blocks에 함수 등록
+with gr.Blocks() as demo:
+    gr.Markdown("# MCP ToolService 예시")
+    gr.Interface(fn=search_similar_chats, inputs=["text", "number"], outputs="json", name="search_similar_chats")
+    gr.Interface(fn=search_chats_by_category, inputs=["text", "number"], outputs="json", name="search_chats_by_category")
+    gr.Interface(fn=search_chats_by_date, inputs=["text", "text", "number"], outputs="json", name="search_chats_by_date")
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(mcp_server=True)