1  import logging
   2  import socket
   3  import subprocess
   4  import psutil
   5  import json
   6  import asyncio
   7  from functools import lru_cache
   8  from ollama import AsyncClient
   9  from modules.tools import available_functions
  10  from modules.system_prompts.system_prompts import SYSTEM_PROMPTS
  11  
  12  logging.basicConfig(level=logging.INFO)
  13  logger = logging.getLogger(__name__)
  14  
  15  @lru_cache(maxsize=1)
  16  def list_models_cached():
  17      return list_models()
  18  
  19  def list_models() -> list:
  20      try:
  21          import requests
  22          response = requests.get('http://localhost:11434/api/tags')
  23          response.raise_for_status()
  24          return [model['name'] for model in response.json().get('models', [])]
  25      except Exception as e:
  26          logger.error(f"Model listing failed: {str(e)}")
  27          return []
  28  
  29  def ensure_ollama_running() -> None:
  30      def is_port_in_use(port):
  31          with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
  32              return s.connect_ex(('localhost', port)) == 0
  33  
  34      def free_port(port):
  35          for proc in psutil.process_iter(['pid', 'name']):
  36              for conn in proc.net_connections(kind='inet'):
  37                  if conn.laddr.port == port:
  38                      try:
  39                          proc.kill()
  40                          return
  41                      except Exception as e:
  42                          logger.error(f"Failed to kill process on port {port}: {e}")
  43  
  44      try:
  45          import requests
  46          response = requests.get("http://localhost:11434")
  47          if response.status_code == 200:
  48              logger.info("Ollama is already running.")
  49              return
  50      except Exception:
  51          logger.info("Ollama not running. Attempting to start it...")
  52          if is_port_in_use(11434):
  53              logger.info("Port 11434 is in use, freeing it...")
  54              free_port(11434)
  55  
  56      try:
  57          subprocess.Popen(["ollama", "serve"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
  58          logger.info("Ollama server started.")
  59      except Exception as e:
  60          logger.error(f"Failed to start Ollama: {e}")
  61  
  62  async def query_ollama(
  63          prompt: list,
  64          model: str,
  65          tools: list,
  66          context: dict,
  67          agent_type: str = 'core_llm'
  68      ) -> dict:
  69          system_prompt = SYSTEM_PROMPTS.get(agent_type.upper(), SYSTEM_PROMPTS['CORE_LLM'])
  70          
  71          formatted_messages = [
  72              {"role": "system", "content": system_prompt},
  73              *prompt,
  74              {"role": "context", "content": str(context)}
  75          ]
  76  
  77          try:
  78              client = AsyncClient()
  79              response = await client.chat(
  80                  model=model,
  81                  messages=formatted_messages,
  82                  tools=tools,
  83                  options={"temperature": 0.4}
  84              )
  85              msg = response.message
  86              assistant_text = msg.content
  87              logger.info("Assistant: %s", msg.tool_calls)
  88              if msg.tool_calls:
  89                  tool_call = msg.tool_calls[0]
  90                  fn_name = tool_call.function.name
  91                  raw_args = tool_call.function.arguments
  92  
  93                  args = raw_args if isinstance(raw_args, dict) else json.loads(raw_args)
  94                  args = {k: int(v) if k == "max_results" and isinstance(v, str) and v.isdigit() else v for k, v in args.items()}
  95  
  96                  if fn_name in available_functions:
  97                      result = await asyncio.to_thread(available_functions[fn_name], **args)
  98                      # messages.append({"role": "tool", "name": fn_name, "content": "\n\n".join(result)})
  99                      mini_sys = {"role": "system",
 100                                   "content": f"Based on this result:\n {result}, try to answer the original question by user: {prompt[-1]['content']}"}
 101  
 102                      followup = await client.chat(model=model, messages=[mini_sys]
 103                                                              , options={"temperature": 0.0})
 104                      return followup.message.content
 105  
 106              return assistant_text
 107          except Exception as e:
 108              logger.exception("Error in query_ollama:")
 109              return f"Ollama Tool Call Error: {str(e)}"
 110