rag-mcp/tools/core.py

"""Core RAG tools: save, retrieve, update, delete."""
import uuid
from typing import Optional
from clients import get_index, get_embedding, merge_with_llm, analyze_relations_with_llm
from utils.logging import get_logger
from utils.validation import validate_content, validate_tag, validate_document_id
from config import AUTO_LINK_THRESHOLD, AUTO_LINK_TOP_K
from .relations import (
    parse_relation,
    format_relation,
    _add_reverse_relation,
    _remove_reverse_relation
)
import time
import threading

logger = get_logger(__name__)

# Rate limiting (in-memory counter)
_rate_limiter = {
    "requests": [],
    "lock": threading.Lock()
}
MAX_REQUESTS_PER_MINUTE = 60

def rate_limit_check() -> bool:
    """
    Check if rate limit is exceeded.

    Returns:
        True if allowed, False if rate limit exceeded
    """
    with _rate_limiter["lock"]:
        now = time.time()
        # Remove requests older than 1 minute
        _rate_limiter["requests"] = [
            req_time for req_time in _rate_limiter["requests"]
            if now - req_time < 60
        ]

        if len(_rate_limiter["requests"]) >= MAX_REQUESTS_PER_MINUTE:
            return False

        _rate_limiter["requests"].append(now)
        return True

def rag_save(content: str, tag: Optional[str] = "general", relations: Optional[str] = None, auto_link: bool = False) -> str:
    """
    Save important information to vector database.

    Args:
        content: Text content to save
        tag: Tag for filtering (default: general)
        relations: Comma-separated relations (e.g., "id1:depends_on,id2:see_also")
        auto_link: Auto-create relations using LLM analysis

    Returns:
        Success message with document ID or error message
    """
    if not rate_limit_check():
        logger.warning("Rate limit exceeded for rag_save")
        return "Error: Rate limit exceeded. Please wait before retrying."

    # Validate inputs
    is_valid, error_msg = validate_content(content)
    if not is_valid:
        logger.warning(f"Content validation failed: {error_msg}")
        return f"Error: {error_msg}"

    is_valid, error_msg = validate_tag(tag)
    if not is_valid:
        logger.warning(f"Tag validation failed: {error_msg}")
        return f"Error: {error_msg}"

    try:
        index = get_index()
        vector = get_embedding(content)
        doc_id = str(uuid.uuid4())

        # Parse manual relations
        rel_list = []
        if relations:
            rel_list = [r.strip() for r in relations.split(',') if r.strip()]

        # Auto-create relations
        auto_relations = []
        if auto_link:
            logger.info(f"Auto-linking document with threshold={AUTO_LINK_THRESHOLD}, top_k={AUTO_LINK_TOP_K}")

            # Search similar documents
            similar_results = index.query(
                vector=vector,
                top_k=AUTO_LINK_TOP_K,
                include_metadata=True
            )

            # Filter by threshold
            similar_docs = []
            for match in similar_results.get("matches", []):
                if match.get("score", 0) >= AUTO_LINK_THRESHOLD:
                    similar_docs.append({
                        "id": match["id"],
                        "text": match["metadata"].get("text", ""),
                        "tag": match["metadata"].get("tag", ""),
                        "score": match["score"]
                    })

            # LLM relation analysis
            if similar_docs:
                analyzed = analyze_relations_with_llm(content, tag, similar_docs)
                for rel in analyzed:
                    rel_str = format_relation(rel["id"], rel["relation"])
                    if rel_str not in rel_list:
                        rel_list.append(rel_str)
                        auto_relations.append(rel_str)

        metadata = {
            "text": content,
            "tag": tag,
            "relations": rel_list
        }

        index.upsert(vectors=[{
            "id": doc_id,
            "values": vector,
            "metadata": metadata
        }])

        # Create bidirectional relations
        for rel_str in rel_list:
            target_id, rel_type = parse_relation(rel_str)
            _add_reverse_relation(target_id, doc_id, rel_type)

        # Format result
        result = f"Saved with ID: {doc_id}"
        if auto_relations:
            result += f"\n\nAuto-linked ({len(auto_relations)}):"
            for rel in auto_relations:
                target_id, rel_type = parse_relation(rel)
                result += f"\n  --[{rel_type}]--> {target_id[:8]}..."
        if relations:
            result += f"\nManual relations: {relations}"

        logger.info(f"Document saved: {doc_id}, tag={tag}, relations={len(rel_list)}")
        return result

    except Exception as e:
        logger.error(f"rag_save failed: {str(e)}", exc_info=True)
        return f"Error: {str(e)}"

def rag_retrieve(query: str, top_k: int = 3, tag: Optional[str] = None) -> str:
    """
    Retrieve relevant information from vector database.

    Args:
        query: Search query
        top_k: Number of results to return (default: 3)
        tag: Filter by specific tag (default: None, search all)

    Returns:
        Formatted search results or error message
    """
    if not rate_limit_check():
        logger.warning("Rate limit exceeded for rag_retrieve")
        return "Error: Rate limit exceeded. Please wait before retrying."

    # Validate query
    is_valid, error_msg = validate_content(query)
    if not is_valid:
        logger.warning(f"Query validation failed: {error_msg}")
        return f"Error: {error_msg}"

    try:
        index = get_index()
        query_vector = get_embedding(query)
        filter_dict = {"tag": {"$eq": tag}} if tag else None

        results = index.query(
            vector=query_vector,
            top_k=top_k,
            include_metadata=True,
            filter=filter_dict
        )

        if not results["matches"]:
            logger.info("No matching documents found")
            return "관련된 정보를 찾지 못했습니다."

        formatted = []
        for i, res in enumerate(results["matches"], 1):
            if "metadata" in res:
                text = res["metadata"]["text"]
                tag_val = res["metadata"].get("tag", "")
                relations = res["metadata"].get("relations", [])
                doc_id = res["id"]
                score = res.get("score", 0)

                entry = f"[{i}] ID: {doc_id} (score: {score:.3f}, tag: {tag_val})"
                if relations:
                    entry += f"\n    Relations: {relations}"
                entry += f"\n    {text}"
                formatted.append(entry)

        logger.info(f"Retrieved {len(formatted)} documents for query")
        return "검색 결과:\n" + "\n---\n".join(formatted)

    except Exception as e:
        logger.error(f"rag_retrieve failed: {str(e)}", exc_info=True)
        return f"Error: {str(e)}"

def rag_update(id: str, new_info: str) -> str:
    """
    Intelligently merge existing and new information.

    Args:
        id: Document ID to update
        new_info: New information to add or merge

    Returns:
        Update summary or error message
    """
    if not rate_limit_check():
        logger.warning("Rate limit exceeded for rag_update")
        return "Error: Rate limit exceeded. Please wait before retrying."

    # Validate inputs
    is_valid, error_msg = validate_document_id(id)
    if not is_valid:
        logger.warning(f"Document ID validation failed: {error_msg}")
        return f"Error: {error_msg}"

    is_valid, error_msg = validate_content(new_info)
    if not is_valid:
        logger.warning(f"New info validation failed: {error_msg}")
        return f"Error: {error_msg}"

    try:
        index = get_index()

        result = index.fetch(ids=[id])
        if id not in result["vectors"]:
            logger.warning(f"Document not found: {id}")
            return f"Error: Not found: {id}"

        old_metadata = result["vectors"][id]["metadata"]
        old_text = old_metadata.get("text", "")
        tag = old_metadata.get("tag", "general")
        relations = old_metadata.get("relations", [])

        merged = merge_with_llm(old_text, new_info)

        vector = get_embedding(merged)
        index.upsert(vectors=[{
            "id": id,
            "values": vector,
            "metadata": {"text": merged, "tag": tag, "relations": relations}
        }])

        logger.info(f"Document updated: {id}")
        return f"Updated: {id}\n\n[기존]\n{old_text}\n\n[새 정보]\n{new_info}\n\n[병합 결과]\n{merged}"

    except Exception as e:
        logger.error(f"rag_update failed: {str(e)}", exc_info=True)
        return f"Error: {str(e)}"

def rag_delete(id: str) -> str:
    """
    Delete document by ID.

    Args:
        id: Document ID to delete

    Returns:
        Success message or error message
    """
    if not rate_limit_check():
        logger.warning("Rate limit exceeded for rag_delete")
        return "Error: Rate limit exceeded. Please wait before retrying."

    # Validate input
    is_valid, error_msg = validate_document_id(id)
    if not is_valid:
        logger.warning(f"Document ID validation failed: {error_msg}")
        return f"Error: {error_msg}"

    try:
        index = get_index()

        # Remove reverse relations
        result = index.fetch(ids=[id])
        if id in result["vectors"]:
            relations = result["vectors"][id]["metadata"].get("relations", [])
            for rel_str in relations:
                target_id, _ = parse_relation(rel_str)
                _remove_reverse_relation(target_id, id)

        index.delete(ids=[id])

        logger.info(f"Document deleted: {id}")
        return f"Deleted: {id}"

    except Exception as e:
        logger.error(f"rag_delete failed: {str(e)}", exc_info=True)
        return f"Error: {str(e)}"