Update server.py

server.py

@@ -1,5 +1,7 @@
 from mcp.server.fastmcp import FastMCP
 from datetime import datetime
+from llama_index import VectorStoreIndex, SimpleDirectoryReader
+from typing import Dict, Optional
 import json
 import os
 import aiohttp  # Necesario para las peticiones HTTP asíncronas
@@ -12,35 +14,57 @@ mcp = FastMCP("OnBase", port=port)
 
 
 
-@mcp.tool()
-def devolver_clave_secreta(
-    clave_ingresada: str,  # Input del usuario
-    clave_secreta: str = "321"  # Clave secreta definida
-) -> str:
+
+# Rutas a los documentos por fuente
+DOCUMENTS_BASE_PATH = "./"
+SOURCES = {
+    "oms": "oms/",
+    #"other": "other/",
+    #"other2": "other2/"
+}
+
+# Cargar índices (se ejecuta una vez al iniciar el servidor)
+indices: Dict[str, VectorStoreIndex] = {}
+for source, rel_path in SOURCES.items():
+    full_path = os.path.join(DOCUMENTS_BASE_PATH, rel_path)
+    if os.path.exists(full_path):
+        documents = SimpleDirectoryReader(full_path).load_data()
+        indices[source] = VectorStoreIndex.from_documents(documents)
+
+# Recurso dinámico para búsqueda
+@mcp.resource(
+    uri="retriever://documentos/{fuente}",
+    name="DocumentRetriever",
+    description="Retrieve documents from differente regulations using semantic search.",
+    mime_type="application/json",
+    tags={"llm", "retrieval"}
+)
+def retrieve_docs(fuente: str, query: str, top_k: int = 3) -> dict:
+    """
+    Parameters:
+        fuente: Source name (for example: oms).
+        query: Search query.
+        top_k: Number of results to return.
     """
-    Verifica si la clave ingresada coincide con la clave secreta.
-    Si coincide, devuelve un JSON con la hora actual y un código.
-    Si no coincide, devuelve un mensaje de error.
+    if fuente not in indices:
+        return {"error": f"Fuente '{fuente}' no disponible. Opciones: {list(indices.keys())}"}
     
-    Args:
-        clave_ingresada: La clave que el usuario ingresa
-        clave_secreta: La clave secreta para comparar (default: "321")
+    retriever = indices[fuente].as_retriever(similarity_top_k=top_k)
+    nodes = retriever.retrieve(query)
     
-    Returns:
-        str: JSON con hora y código o mensaje de error
-    """
-    if clave_ingresada == clave_secreta:
-        respuesta = {
-            "hora_actual": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            "codigo": "codigo 123",
-            "status": "success"
+    results = [
+        {
+            "content": node.get_content(),
+            "metadata": node.metadata,
+            "score": node.score
         }
-        return json.dumps(respuesta, ensure_ascii=False)
-    else:
-        return json.dumps({
-            "error": "Clave incorrecta",
-            "status": "fail"
-        }, ensure_ascii=False)
+        for node in nodes
+    ]
+    
+    return {"results": results}
+
+
+
     
 
 @mcp.tool()