Added 'ProviderError' for error handling

error handling i.e not showing a traceback on the ai-ops-cli client

src/agent/llm/__init__.py

@@ -1,4 +1,12 @@
-from src.agent.llm.llm import LLM, Provider, Ollama, OpenRouter
+"""Exposes implemented LLM functionalities"""
+
+from src.agent.llm.llm import (
+    LLM,
+    Provider,
+    Ollama,
+    OpenRouter,
+    ProviderError
+)
 
 AVAILABLE_PROVIDERS = {
     'ollama': {'class': Ollama, 'key_required': False},

src/agent/llm/llm.py

@@ -52,6 +52,11 @@ def query(self, messages: list):
         """Implement to makes query to the LLM provider"""
 
 
+class ProviderError(Exception):
+    """Just a wrapper to Exception for error handling
+    when an error is caused by the LLM provider"""
+
+
 @dataclass
 class Ollama(Provider):
     """Ollama Interface"""
@@ -74,7 +79,7 @@ def query(self, messages: list):
             for chunk in stream:
                 yield chunk['message']['content']
         except ResponseError as err:
-            raise RuntimeError(err)
+            raise ProviderError(err)
 
 
 @dataclass

src/api.py

@@ -21,6 +21,7 @@
 - /collections/new: Creates a new Collection.
 """
 import os
+import json
 
 from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException, Body
@@ -33,6 +34,7 @@
 # from src.agent.knowledge import Store
 from src.agent.plan import TaskStatus
 from src.agent.tools import TOOLS
+from src.agent.llm import ProviderError
 
 load_dotenv()
 
@@ -170,9 +172,12 @@ def delete_session(sid: int):
 # --- AGENT RELATED
 
 def query_generator(sid: int, q: str):
-    stream = agent.query(sid, q, rag=False)
-    for chunk in stream:
-        yield chunk
+    try:
+        stream = agent.query(sid, q, rag=False)
+        for chunk in stream:
+            yield chunk
+    except ProviderError as err:
+        yield json.dumps({'error': str(err)})
 
 
 @app.post('/session/{sid}/query/')