[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

EdgeCraftRAG/edgecraftrag/api/v1/agent.py

@@ -4,11 +4,12 @@
 import json
 import os
 import time
+
 from edgecraftrag.api_schema import AgentCreateIn
 from edgecraftrag.base import AgentType
+from edgecraftrag.config_repository import MilvusConfigRepository, save_agent_configurations
 from edgecraftrag.context import ctx
 from edgecraftrag.env import AGENT_FILE
-from edgecraftrag.config_repository import MilvusConfigRepository, save_agent_configurations
 from fastapi import FastAPI, HTTPException, status
 
 agent_app = FastAPI()
@@ -21,14 +22,16 @@ async def get_all_agents():
     agents = ctx.get_agent_mgr().get_agents()
     active_id = ctx.get_agent_mgr().get_active_agent_id()
     for k, agent in agents.items():
-        out.append(AgentCreateIn(
-            idx=agent.idx,
-            name=agent.name,
-            type=agent.comp_subtype,
-            pipeline_idx=agent.pipeline_idx,
-            configs=agent.configs,
-            active=True if agent.idx == active_id else False
-        ))
+        out.append(
+            AgentCreateIn(
+                idx=agent.idx,
+                name=agent.name,
+                type=agent.comp_subtype,
+                pipeline_idx=agent.pipeline_idx,
+                configs=agent.configs,
+                active=True if agent.idx == active_id else False,
+            )
+        )
     return out
 
 
@@ -44,7 +47,7 @@ async def get_agent(name):
             type=agent.comp_subtype,
             pipeline_idx=agent.pipeline_idx,
             configs=agent.configs,
-            active=isactive
+            active=isactive,
         )
     else:
         raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)
@@ -95,7 +98,7 @@ async def delete_agent(name):
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e))
 
 
-# GET Agent Type defualt configs
+# GET Agent Type default configs
 @agent_app.get(path="/v1/settings/agents/configs/{agent_type}")
 async def get_agent_default_configs(agent_type):
     try:

EdgeCraftRAG/edgecraftrag/api/v1/chatqna.py

@@ -1,15 +1,16 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-import requests
-import json
 import asyncio
-from typing import List
+import json
 from concurrent.futures import ThreadPoolExecutor
+from typing import List
+
+import requests
 from comps.cores.proto.api_protocol import ChatCompletionRequest
 from edgecraftrag.api_schema import RagOut
 from edgecraftrag.context import ctx
-from edgecraftrag.utils import serialize_contexts, stream_generator, chain_async_generators
+from edgecraftrag.utils import chain_async_generators, serialize_contexts, stream_generator
 from fastapi import Body, FastAPI, HTTPException, status
 from fastapi.responses import StreamingResponse
 
@@ -25,7 +26,10 @@ async def retrieval(request: ChatCompletionRequest):
         if active_kb:
             request.user = active_kb
         else:
-            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Retrieval needs to have an active knowledgebase")
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Retrieval needs to have an active knowledgebase",
+            )
         contexts = await ctx.get_pipeline_mgr().run_retrieve_postprocess(chat_request=request)
         serialized_contexts = serialize_contexts(contexts)
 
@@ -63,16 +67,10 @@ async def chatqna(request: ChatCompletionRequest):
                 request.model = generator.model_id
 
         if request.stream:
-            run_pipeline_gen, contexts = await ctx.get_pipeline_mgr().run_pipeline(
-                chat_request=request
-            )
-            return StreamingResponse(
-                save_session(sessionid, run_pipeline_gen), media_type="text/plain"
-            )
+            run_pipeline_gen, contexts = await ctx.get_pipeline_mgr().run_pipeline(chat_request=request)
+            return StreamingResponse(save_session(sessionid, run_pipeline_gen), media_type="text/plain")
         else:
-            ret, contexts = await ctx.get_pipeline_mgr().run_pipeline(
-                chat_request=request
-            )
+            ret, contexts = await ctx.get_pipeline_mgr().run_pipeline(chat_request=request)
             ctx.get_session_mgr().save_current_message(sessionid, "assistant", str(ret))
             return str(ret)
 
@@ -105,11 +103,11 @@ async def res_gen_json():
                     yield token.replace("\n", "\\n")
 
             # Reconstruct RagOut in stream response
-            query_gen = stream_generator("{\"query\":\"" + request.messages + "\",")
+            query_gen = stream_generator('{"query":"' + request.messages + '",')
 
             s_contexts = json.dumps(serialize_contexts(contexts))
-            context_gen = stream_generator("\"contexts\":" + s_contexts + ",\"response\":\"")
-            final_gen = stream_generator("\"}")
+            context_gen = stream_generator('"contexts":' + s_contexts + ',"response":"')
+            final_gen = stream_generator('"}')
             output_gen = chain_async_generators([query_gen, context_gen, res_gen_json(), final_gen])
 
             return StreamingResponse(output_gen, media_type="text/plain")

EdgeCraftRAG/edgecraftrag/api/v1/data.py

@@ -1,15 +1,15 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+import json
 import os
 from typing import List
 
 from edgecraftrag.api_schema import DataIn, FilesIn
-from edgecraftrag.context import ctx
-from fastapi import FastAPI, File, HTTPException, UploadFile, status
-import json
 from edgecraftrag.config_repository import MilvusConfigRepository
+from edgecraftrag.context import ctx
 from edgecraftrag.env import UI_DIRECTORY
+from fastapi import FastAPI, File, HTTPException, UploadFile, status
 
 data_app = FastAPI()
 
@@ -52,22 +52,22 @@ async def redindex_data():
 # Gets the current nodelist
 @data_app.get(path="/v1/data/nodes")
 async def get_nodes_with_kb(kb_name=None):
-    node_lists  = {}
+    node_lists = {}
     active_pl = ctx.get_pipeline_mgr().get_active_pipeline()
     if kb_name:
         kb = ctx.get_knowledge_mgr().get_knowledge_base_by_name_or_id(kb_name)
     else:
         kb = ctx.get_knowledge_mgr().get_active_knowledge_base()
     if active_pl.indexer.comp_subtype == "faiss_vector":
-        return  active_pl.indexer.docstore.docs
+        return active_pl.indexer.docstore.docs
     elif active_pl.indexer.comp_subtype == "milvus_vector":
         collection_name = kb.name + active_pl.name
-        Milvus_node_list = MilvusConfigRepository.create_connection(collection_name,1, active_pl.indexer.vector_url)
+        Milvus_node_list = MilvusConfigRepository.create_connection(collection_name, 1, active_pl.indexer.vector_url)
         results = Milvus_node_list.get_configs(output_fields=["text", "_node_content", "doc_id"])
         for node_list in results:
             text = node_list.get("text")
-            node_content =  json.loads(node_list.get("_node_content"))
-            node_content["doc_id"]=node_list.get("doc_id")
+            node_content = json.loads(node_list.get("_node_content"))
+            node_content["doc_id"] = node_list.get("doc_id")
             node_content["text"] = text
             node_lists[node_content.get("id_")] = node_content
         return node_lists
@@ -81,12 +81,10 @@ async def get_nodes_by_document_name(document_name: str):
     all_nodes = await get_nodes_with_kb()
     matching_nodes = []
     for node in all_nodes.values() if isinstance(all_nodes, dict) else all_nodes:
-        metadata = node.get('metadata', {}) if isinstance(node, dict) else getattr(node, 'metadata', {})
-        node_file_name = metadata.get('file_name', '')
-        node_file_path = metadata.get('file_path', '')
-        if (node_file_name == document_name or
-            document_name in node_file_name or
-            document_name in node_file_path):
+        metadata = node.get("metadata", {}) if isinstance(node, dict) else getattr(node, "metadata", {})
+        node_file_name = metadata.get("file_name", "")
+        node_file_path = metadata.get("file_path", "")
+        if node_file_name == document_name or document_name in node_file_name or document_name in node_file_path:
             matching_nodes.append(node)
     return matching_nodes
 
@@ -100,23 +98,20 @@ async def get_document_names():
 
     documents = {}
     for node in all_nodes.values() if isinstance(all_nodes, dict) else all_nodes:
-        metadata = node.get('metadata', {}) if isinstance(node, dict) else getattr(node, 'metadata', {})
-        file_name = metadata.get('file_name')
-        file_path = metadata.get('file_path')
+        metadata = node.get("metadata", {}) if isinstance(node, dict) else getattr(node, "metadata", {})
+        file_name = metadata.get("file_name")
+        file_path = metadata.get("file_path")
         if file_name and file_name not in documents:
             documents[file_name] = {
                 "file_name": file_name,
                 "file_path": file_path,
-                "file_type": metadata.get('file_type', 'unknown'),
-                "chunk_count": 0
+                "file_type": metadata.get("file_type", "unknown"),
+                "chunk_count": 0,
             }
         if file_name:
             documents[file_name]["chunk_count"] += 1
 
-    return {
-        "total_documents": len(documents),
-        "documents": list(documents.values())
-    }
+    return {"total_documents": len(documents), "documents": list(documents.values())}
 
 
 # Upload files by a list of file_path
@@ -145,8 +140,7 @@ async def get_files():
 async def get_nodes_by_document_name(document_name: str) -> List[dict]:
     pl = ctx.get_pipeline_mgr().get_active_pipeline()
     if pl is None:
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND, detail="No active pipeline")
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="No active pipeline")
 
     nodelist = ctx.get_node_mgr().get_nodes(pl.node_parser.idx)
 
@@ -155,13 +149,11 @@ async def get_nodes_by_document_name(document_name: str) -> List[dict]:
 
     matching_nodes = []
     for node in nodelist:
-        if hasattr(node, 'metadata') and node.metadata:
-            node_file_name = node.metadata.get('file_name', '')
-            node_file_path = node.metadata.get('file_path', '')
+        if hasattr(node, "metadata") and node.metadata:
+            node_file_name = node.metadata.get("file_name", "")
+            node_file_path = node.metadata.get("file_path", "")
 
-            if (node_file_name == document_name or
-                document_name in node_file_name or
-                    document_name in node_file_path):
+            if node_file_name == document_name or document_name in node_file_name or document_name in node_file_path:
                 node_dict = node.model_dump()
                 matching_nodes.append(node_dict)
 
@@ -173,8 +165,7 @@ async def get_nodes_by_document_name(document_name: str) -> List[dict]:
 async def get_node_by_id(node_id: str) -> dict:
     pl = ctx.get_pipeline_mgr().get_active_pipeline()
     if pl is None:
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND, detail="No active pipeline")
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="No active pipeline")
 
     nodelist = ctx.get_node_mgr().get_nodes(pl.node_parser.idx)
 
@@ -193,8 +184,7 @@ async def get_node_by_id(node_id: str) -> dict:
 async def get_document_names():
     pl = ctx.get_pipeline_mgr().get_active_pipeline()
     if pl is None:
-        raise HTTPException(
-            status_code=status.HTTP_404_NOT_FOUND, detail="No active pipeline")
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="No active pipeline")
 
     nodelist = ctx.get_node_mgr().get_nodes(pl.node_parser.idx)
 
@@ -203,25 +193,22 @@ async def get_document_names():
 
     documents = {}
     for node in nodelist:
-        if hasattr(node, 'metadata') and node.metadata:
-            file_name = node.metadata.get('file_name')
-            file_path = node.metadata.get('file_path')
+        if hasattr(node, "metadata") and node.metadata:
+            file_name = node.metadata.get("file_name")
+            file_path = node.metadata.get("file_path")
 
             if file_name and file_name not in documents:
                 documents[file_name] = {
                     "file_name": file_name,
                     "file_path": file_path,
-                    "file_type": node.metadata.get('file_type', 'unknown'),
-                    "chunk_count": 0
+                    "file_type": node.metadata.get("file_type", "unknown"),
+                    "chunk_count": 0,
                 }
 
             if file_name:
                 documents[file_name]["chunk_count"] += 1
 
-    return {
-        "total_documents": len(documents),
-        "documents": list(documents.values())
-    }
+    return {"total_documents": len(documents), "documents": list(documents.values())}
 
 
 # GET a file