Update app.py

app.py

@@ -1,11 +1,14 @@
 import os
 import subprocess
-import threading
-import time
 import asyncio
-import uvicorn
+import time
+import threading
+import gc
+import psutil
+import torch
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+import uvicorn
 from langchain.prompts import PromptTemplate
 from langchain_community.chat_models import ChatOllama
 from langchain_community.tools import DuckDuckGoSearchRun
@@ -13,57 +16,54 @@ from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
 from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
 from langgraph.graph import END, StateGraph
 from typing_extensions import TypedDict
-from fastapi.responses import StreamingResponse
-from ollama import AsyncClient
-import gc
-import psutil
-import torch
-from functools import lru_cache
 
 # Asegúrate de tener la librería Ollama instalada
 # pip install ollama langchain langchain_community langgraph
 
-# Ruta de Ollama
+# Configuración de Ollama y su servicio
 OLLAMA = os.path.expanduser("~/ollama")
-if not os.path.exists(OLLAMA):
-    print("Ollama no encontrado, descargando...")
-    subprocess.run("curl -L https://ollama.com/download/ollama-linux-amd64 -o ~/ollama", shell=True)
-    os.chmod(OLLAMA, 0o755)
 
-# Iniciar el servidor Ollama en segundo plano
-async def ollama_service_thread():
+def download_ollama_model(model_name='hf.co/MaziyarPanahi/Llama-3.2-3B-Instruct-uncensored-GGUF:IQ1_S'):
+    try:
+        if not os.path.exists(OLLAMA):
+            print("Ollama no encontrado, descargando...")
+            subprocess.run("curl -L https://ollama.com/download/ollama-linux-amd64 -o ~/ollama", shell=True)
+            os.chmod(OLLAMA, 0o755)
+        
+        print(f"Descargando el modelo: {model_name}")
+        subprocess.run(["~/ollama", "pull", model_name], check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error al descargar el modelo: {e}")
+        raise
+
+# Función asíncrona para manejar la descarga del modelo
+async def async_download_ollama_model():
+    await asyncio.to_thread(download_ollama_model)
+
+# Iniciar el servidor Ollama en un hilo
+def ollama_service_thread():
     print("Iniciando el servicio de Ollama")
-    subprocess.run("~/ollama serve -1", shell=True)  # -1 indica mantener el servidor en ejecución indefinidamente
+    subprocess.run("~/ollama serve", shell=True)
 
-# Iniciar el hilo de servicio Ollama
-print("Creando y comenzando el hilo del servicio Ollama")
-OLLAMA_SERVICE_THREAD = threading.Thread(target=asyncio.run, args=(ollama_service_thread(),))
+# Crear un hilo para iniciar Ollama
+OLLAMA_SERVICE_THREAD = threading.Thread(target=ollama_service_thread)
 OLLAMA_SERVICE_THREAD.start()
 
-# Esperar a que Ollama se inicie
+# Esperar a que Ollama esté listo
 print("Esperando a que Ollama inicie...")
 time.sleep(10)
 
 # Descargar el modelo de Hugging Face si no está disponible
-async def download_ollama_model(model_name='hf.co/MaziyarPanahi/Llama-3.2-3B-Instruct-uncensored-GGUF:IQ1_S'):
-    try:
-        print(f"Descargando el modelo: {model_name}")
-        subprocess.run(["ollama", "pull", model_name], check=True)
-    except subprocess.CalledProcessError as e:
-        print(f"Error al descargar el modelo: {e}")
-        raise
-
-# Descargar el modelo de Ollama en el hilo principal
-download_ollama_model("hf.co/MaziyarPanahi/Llama-3.2-3B-Instruct-uncensored-GGUF:IQ1_S")
+asyncio.run(async_download_ollama_model())
 
-# Crear la aplicación FastAPI
+# Crear instancia de FastAPI
 app = FastAPI()
 
-# Definir el modelo de datos para recibir las consultas en la API
+# Definición del modelo de datos para recibir las consultas en la API
 class QueryRequest(BaseModel):
     query: str
 
-# Definir el modelo de lenguaje de Ollama
+# Definir el modelo de lenguaje de Ollama (sin 'temperature')
 local_llm = 'hf.co/MaziyarPanahi/Llama-3.2-3B-Instruct-uncensored-GGUF:IQ1_S'
 llama3 = ChatOllama(model=local_llm)
 
@@ -71,12 +71,7 @@ llama3 = ChatOllama(model=local_llm)
 wrapper = DuckDuckGoSearchAPIWrapper(max_results=1)
 web_search_tool = DuckDuckGoSearchRun(api_wrapper=wrapper)
 
-# Implementar optimización extrema mediante LRU Cache
-@lru_cache(maxsize=1024)  # Cache de tamaño limitado para resultados previos
-async def cached_search(query):
-    return await web_search_tool.invoke(query)
-
-# Definir los prompts para generación y enrutamiento
+# Definición de los prompts para generación y enrutamiento
 generate_prompt = PromptTemplate(
     template="""
     <|begin_of_text|>
@@ -139,26 +134,25 @@ class GraphState(TypedDict):
     context: str
 
 # Nodos de procesamiento
-async def generate(state):
+def generate(state):
     print("Step: Generating Final Response")
     question = state["question"]
     context = state["context"]
-    generation = await generate_chain.invoke({"context": context, "question": question})
+    generation = generate_chain.invoke({"context": context, "question": question})
     return {"generation": generation}
 
-async def transform_query(state):
+def transform_query(state):
     print("Step: Optimizing Query for Web Search")
     question = state['question']
-    gen_query = await query_chain.invoke({"question": question})
+    gen_query = query_chain.invoke({"question": question})
     search_query = gen_query.get("query", "")  # Asegurarnos de que estamos obteniendo la clave correcta
     return {"search_query": search_query}
 
-async def web_search(state):
+def web_search(state):
     search_query = state['search_query']
     print(f'Step: Searching the Web for: "{search_query}"')
     try:
-        # Se usa la caché para optimizar los resultados
-        search_result = await cached_search(search_query)
+        search_result = web_search_tool.invoke(search_query)
         if isinstance(search_result, str):  # Si la respuesta es una cadena, la convertimos en un diccionario
             print(f"Respuesta de búsqueda web es cadena: {search_result}")
             return {"context": search_result}
@@ -170,10 +164,10 @@ async def web_search(state):
         print(f"Web search failed: {e}")
         return None  # Si la búsqueda falla, no devuelve contexto
 
-async def route_question(state):
+def route_question(state):
     print("Step: Routing Query")
     question = state['question']
-    output = await question_router.invoke({"question": question})
+    output = question_router.invoke({"question": question})
     if output.get('choice') == "web_search":
         print("Step: Routing Query to Web Search")
         return "websearch"
@@ -194,8 +188,28 @@ workflow.set_conditional_entry_point(
         "generate": "generate",
     },
 )
+workflow.add_edge("transform_query", "websearch")
+workflow.add_edge("websearch", "generate")
+workflow.add_edge("generate", END)
+
+# Compilar el agente
+local_agent = workflow.compile()
+
+# Función para ejecutar el agente
+def run_agent_parallel(query):
+    output = local_agent.invoke({"question": query})
+    if "generation" not in output:  # Si la búsqueda web falló y no hubo respuesta de generación
+        print("Web search failed, using Ollama model directly.")
+        return generate({"question": query, "context": ""})["generation"]  # Generar directamente
+    return output['generation']
+
+# Lógica del servidor FastAPI
+@app.post("/query")
+async def query_endpoint(request: QueryRequest):
+    query = request.query
+    return {"response": run_agent_parallel(query)}
 
-# Gestión de recursos de CPU, RAM y GPU
+# Lógica de recursos
 def release_resources():
     try:
         torch.cuda.empty_cache()
@@ -204,16 +218,16 @@ def release_resources():
         print(f"Failed to release resources: {e}")
 
 def resource_manager():
-    MAX_RAM_PERCENT = 1  # Ajustar según sea necesario
+    MAX_RAM_PERCENT = 1
     MAX_CPU_PERCENT = 1
     MAX_GPU_PERCENT = 1
-    MAX_RAM_MB = 1  # Ajustar según la memoria disponible
+    MAX_RAM_MB = 1
 
     while True:
         try:
             virtual_mem = psutil.virtual_memory()
             current_ram_percent = virtual_mem.percent
-            current_ram_mb = virtual_mem.used / (1 * 1)  # Convertir a MB
+            current_ram_mb = virtual_mem.used / (1 * 1)  # Convert to MB
 
             if current_ram_percent > MAX_RAM_PERCENT or current_ram_mb > MAX_RAM_MB:
                 release_resources()
@@ -224,35 +238,15 @@ def resource_manager():
 
             if torch.cuda.is_available():
                 gpu = torch.cuda.current_device()
-                gpu_mem = torch.cuda.memory_allocated(gpu) / (1024 * 1024)  # Convertir a MB
+                gpu_mem = torch.cuda.memory_percent(gpu)
 
                 if gpu_mem > MAX_GPU_PERCENT:
                     release_resources()
 
         except Exception as e:
-            print(f"Error en el gestor de recursos: {e}")
+            print(f"Error in resource manager: {e}")
 
 resource_manager()
 
-# Procesar la consulta en paralelo
-async def process_query_in_parallel(query):
-    try:
-        state = GraphState(question=query, generation="", search_query="", context="")
-        return await workflow.invoke(state)
-    except Exception as e:
-        print(f"Error en la ejecución paralela: {e}")
-        raise
-
-# Ruta de la API para manejar consultas
-@app.post("/query")
-async def query_handler(request: QueryRequest):
-    try:
-        query = request.query
-        result = await process_query_in_parallel(query)  # Llamada asíncrona
-        return {"results": result}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-# Ejecutar el servidor FastAPI
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)