pylint check

src/agent/agent.py

@@ -111,7 +111,7 @@ def extract_plan(self, plan_nl):
             if json_match:
                 plan_data = json.loads(json_match.group())
             else:
-                print(f'PlanError: Response: \n{response["message"]["content"]}')
+                print(f'PlanError:\n{response["message"]["content"]}')
                 return None
 
         tasks = []
@@ -134,9 +134,9 @@ def execute_plan(self, sid):
             return None
 
         msg = messages[-1] if messages[-1].role == Role.ASSISTANT else messages[-2]
+
         plan = self.extract_plan(msg.content)
-        for tasks_status in plan.execute():
-            yield tasks_status
+        yield from plan.execute()
 
         self.mem.store_plan(sid, plan)
 

src/agent/llm.py

@@ -39,6 +39,7 @@ class LLM:
     """Ollama model interface"""
     model: str
     client_url: str = 'http://localhost:11434'
+    client: Client | None = None
 
     def __post_init__(self):
         if self.model not in AVAILABLE_MODELS.keys():

src/agent/memory/base.py

@@ -25,12 +25,11 @@ class Role(StrEnum):
     def from_str(item):
         if item == 'user':
             return Role.USER
-        elif item == 'assistant':
+        if item == 'assistant':
             return Role.ASSISTANT
-        elif item == 'system':
+        if item == 'system':
             return Role.SYS
-        else:
-            return None
+        return None
 
 
 @dataclass
@@ -88,7 +87,7 @@ def from_json(path: str):
 
             plans = None
             if 'plans' in data and data['plans'] is not None:
-                plans = [plan_list for plan_list in data['plans']]
+                plans = list(data['plans'])
 
             session = Session(
                 name=data['name'],
@@ -150,12 +149,15 @@ def save_session(self, sid: int):
 
         session = self.sessions[sid]
         self.delete_session(sid)
-        with open(f'{SESSIONS_PATH}/{sid}__{session.name}.json', 'w+', encoding='utf-8') as fp:
+
+        path = f'{SESSIONS_PATH}/{sid}__{session.name}.json'
+        plans = session.plans_to_dict_list() if session.plans else None
+        with open(path, 'w+', encoding='utf-8') as fp:
             data = {
                 'id': sid,
                 'name': session.name,
                 'messages': session.messages_to_dict_list(),
-                'plans': session.plans_to_dict_list() if session.plans is not None else None
+                'plans': plans
             }
             json.dump(data, fp)
 

src/agent/tools/__init__.py

@@ -1,3 +1,4 @@
+"""Tool package"""
 from pathlib import Path
 
 from src.agent.tools.base import Tool
@@ -11,4 +12,3 @@
 for path in TOOLS_PATH.iterdir():
     tool = Tool.load_tool(str(path))
     TOOLS.append(tool)
-

src/agent/tools/base.py

@@ -48,4 +48,3 @@ def get_documentation(self):
 Arguments:
     {self.args_description}          
 """
-

src/api.py

@@ -180,7 +180,7 @@ def query(sid: int, body: dict = Body(...)):
     """
     q = body.get("query")
     if not q:
-        raise HTTPException(status_code=400, detail="Query parameter is required")
+        raise HTTPException(status_code=400, detail="Query parameter required")
     return StreamingResponse(query_generator(sid, q))
 
 
@@ -209,7 +209,7 @@ def execute_plan_stream(sid: int):
     """Generator for plan execution and status updates"""
     execution = agent.execute_plan(sid)
     for iteration in execution:
-        for i, task in enumerate(iteration):
+        for task in iteration:
             if task.status == TaskStatus.DONE:
                 task_str = f'ai-ops:~$ {task.command}\n{task.output}\n'
                 yield task_str
@@ -219,8 +219,8 @@ def execute_plan_stream(sid: int):
     for p in plan.plan_to_dict_list():
         eval_results += f'{p["command"]}\n{p["output"]}\n\n'
 
-    for chunk in query_generator(sid, eval_results):
-        yield chunk
+    yield from query_generator(sid, eval_results)
+
 
 @app.get('/session/{sid}/plan/execute')
 def execute_plan(sid: int):