flowllm 0.1.2__py3-none-any.whl → 0.1.5__py3-none-any.whl

flowllm/service/base_service.py

@@ -1,4 +1,4 @@
-from abc import abstractmethod, ABC
+from abc import ABC
 from typing import Dict, Optional
 
 from loguru import logger
@@ -6,16 +6,19 @@ from pydantic import create_model, Field
 
 from flowllm.config.pydantic_config_parser import PydanticConfigParser
 from flowllm.context.service_context import C
+from flowllm.flow.base_tool_flow import BaseToolFlow
 from flowllm.schema.flow_request import FlowRequest
 from flowllm.schema.service_config import ServiceConfig
 from flowllm.schema.tool_call import ParamAttrs
-from flowllm.utils.common_utils import snake_to_camel
+from flowllm.utils.common_utils import snake_to_camel, print_banner
 
 
 class BaseService(ABC):
     TYPE_MAPPING = {
         "str": str,
+        "string": str,
         "int": int,
+        "integer": int,
         "float": float,
         "bool": bool,
         "list": list,
@@ -24,10 +27,16 @@ class BaseService(ABC):
 
     def __init__(self, service_config: ServiceConfig):
         self.service_config = service_config
-
         self.mcp_config = self.service_config.mcp
         self.http_config = self.service_config.http
+
+    def __enter__(self):
+        C.prepare_sse_mcp()
         C.init_by_service_config(self.service_config)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        C.stop_by_service_config()
 
     @classmethod
     def get_service(cls, *args, parser: type[PydanticConfigParser] = PydanticConfigParser) -> "BaseService":
@@ -49,20 +58,28 @@ class BaseService(ABC):
 
         return create_model(f"{snake_to_camel(flow_name)}Model", __base__=FlowRequest, **fields)
 
-    def integrate_tool_flow(self, tool_flow_name: str):
+    def integrate_flow(self, tool_flow: BaseToolFlow):
+        ...
+
+    def integrate_stream_flow(self, tool_flow: BaseToolFlow):
         ...
 
-    def integrate_tool_flows(self):
+    def integrate_flows(self):
         for tool_flow_name in C.tool_flow_names:
-            self.integrate_tool_flow(tool_flow_name)
-            logger.info(f"integrate flow_endpoint={tool_flow_name}")
+            tool_flow: BaseToolFlow = C.get_tool_flow(tool_flow_name)
+            if tool_flow.stream:
+                self.integrate_stream_flow(tool_flow)
+                logger.info(f"integrate stream_endpoint={tool_flow_name}")
 
-    def __enter__(self):
-        return self
+            else:
+                self.integrate_flow(tool_flow)
+                logger.info(f"integrate endpoint={tool_flow_name}")
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        ...
+    def __call__(self, logo: str = ""):
+        self.integrate_flows()
+        if logo:
+            print_banner(name=logo, service_config=self.service_config, width=400)
+        self.execute()
 
-    @abstractmethod
-    def __call__(self):
-        ...
+    def execute(self):
+        ...

flowllm/service/http_service.py

@@ -1,14 +1,16 @@
 import asyncio
-from functools import partial
+from typing import AsyncGenerator
 
 import uvicorn
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
 from loguru import logger
 
 from flowllm.context.service_context import C
 from flowllm.flow.base_tool_flow import BaseToolFlow
 from flowllm.schema.flow_response import FlowResponse
+from flowllm.schema.flow_stream_chunk import FlowStreamChunk
 from flowllm.service.base_service import BaseService
 
 
@@ -18,60 +20,67 @@ class HttpService(BaseService):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.app = FastAPI(title="FlowLLM", description="HTTP API for FlowLLM")
+        self.app.add_middleware(CORSMiddleware,
+                                allow_origins=["*"],
+                                allow_credentials=True,
+                                allow_methods=["*"],
+                                allow_headers=["*"])
 
-        # Add CORS middleware
-        self.app.add_middleware(
-            CORSMiddleware,
-            allow_origins=["*"],
-            allow_credentials=True,
-            allow_methods=["*"],
-            allow_headers=["*"],
-        )
-
-        # Add health check endpoint
         self.app.get("/health")(self.health_check)
 
     @staticmethod
     def health_check():
         return {"status": "healthy"}
 
-    def integrate_tool_flow(self, tool_flow_name: str):
-        tool_flow: BaseToolFlow = C.get_tool_flow(tool_flow_name)
-        request_model = self._create_pydantic_model(tool_flow_name, tool_flow.tool_call.input_schema)
+    def integrate_flow(self, tool_flow: BaseToolFlow):
+        if "http" not in tool_flow.service_type:
+            return
+
+        request_model = self._create_pydantic_model(tool_flow.name, tool_flow.tool_call.input_schema)
 
         async def execute_endpoint(request: request_model) -> FlowResponse:
-            loop = asyncio.get_event_loop()
-            response: FlowResponse = await loop.run_in_executor(
-                executor=C.thread_pool,
-                func=partial(tool_flow.__call__, **request.model_dump()))  # noqa
+            return await tool_flow(**request.model_dump())
 
-            return response
+        self.app.post(f"/{tool_flow.name}", response_model=FlowResponse)(execute_endpoint)
 
-        endpoint_path = f"/{tool_flow.name}"
-        self.app.post(endpoint_path, response_model=FlowResponse)(execute_endpoint)
+    @staticmethod
+    def gen_stream_response(queue: asyncio.Queue):
+        async def generate_stream() -> AsyncGenerator[bytes, None]:
+            while True:
+                stream_chunk: FlowStreamChunk = await queue.get()
+                if stream_chunk.done:
+                    yield f"data:[DONE]\n\n".encode('utf-8')
+                    break
+                else:
+                    yield f"data:{stream_chunk.model_dump_json()}\n\n".encode("utf-8")
 
-    def integrate_tool_flows(self):
-        super().integrate_tool_flows()
+        return StreamingResponse(generate_stream(), media_type="text/event-stream")
 
-        async def execute_endpoint() -> list:
-            loop = asyncio.get_event_loop()
+    def integrate_stream_flow(self, tool_flow: BaseToolFlow):
+        if "http" not in tool_flow.service_type:
+            return
 
-            def list_tool_flows() -> list:
-                tool_flow_schemas = []
-                for name, tool_flow in C.tool_flow_dict.items():
-                    assert isinstance(tool_flow, BaseToolFlow)
-                    tool_flow_schemas.append(tool_flow.tool_call.simple_input_dump())
-                return tool_flow_schemas
+        request_model = self._create_pydantic_model(tool_flow.name, tool_flow.tool_call.input_schema)
 
-            return await loop.run_in_executor(executor=C.thread_pool, func=list_tool_flows)  # noqa
+        async def execute_stream_endpoint(request: request_model) -> StreamingResponse:
+            stream_queue = asyncio.Queue()
+            asyncio.create_task(tool_flow(stream_queue=stream_queue, **request.model_dump()))
+            return self.gen_stream_response(stream_queue)
 
-        endpoint_path = "/list"
-        self.app.get(endpoint_path, response_model=list)(execute_endpoint)
-        logger.info(f"integrate endpoint={endpoint_path}")
+        self.app.post(f"/{tool_flow.name}")(execute_stream_endpoint)
+
+    def integrate_flows(self):
+        super().integrate_flows()
+
+        async def execute_endpoint() -> list:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(executor=C.thread_pool, func=C.list_flow_schemas)  # noqa
 
-    def __call__(self):
-        self.integrate_tool_flows()
+        endpoint_path = "list"
+        self.app.get(f"/{endpoint_path}", response_model=list)(execute_endpoint)
+        logger.info(f"integrate endpoint={endpoint_path}")
 
+    def execute(self):
         uvicorn.run(self.app,
                     host=self.http_config.host,
                     port=self.http_config.port,

flowllm/service/mcp_service.py

@@ -1,12 +1,9 @@
-import asyncio
-from functools import partial
-
 from fastmcp import FastMCP
 from fastmcp.tools import FunctionTool
 
-from flowllm.context.service_context import C
-from flowllm.flow.base_tool_flow import BaseToolFlow
-from flowllm.service.base_service import BaseService
+from .base_service import BaseService
+from ..context.service_context import C
+from ..flow.base_tool_flow import BaseToolFlow
 
 
 @C.register_service("mcp")
@@ -16,32 +13,29 @@ class MCPService(BaseService):
         super().__init__(*args, **kwargs)
         self.mcp = FastMCP(name="FlowLLM")
 
-    def integrate_tool_flow(self, tool_flow_name: str):
-        tool_flow: BaseToolFlow = C.get_tool_flow(tool_flow_name)
-        request_model = self._create_pydantic_model(tool_flow_name, tool_flow.tool_call.input_schema)
+    def integrate_flow(self, tool_flow: BaseToolFlow):
+        if "mcp" not in tool_flow.service_type:
+            return
+
+        request_model = self._create_pydantic_model(tool_flow.name, tool_flow.tool_call.input_schema)
 
-        async def execute_flow_async(**kwargs) -> str:
-            request: request_model = request_model(**kwargs)
-            loop = asyncio.get_event_loop()
-            response = await loop.run_in_executor(
-                executor=C.thread_pool,
-                func=partial(tool_flow.__call__, **request.model_dump()))  # noqa
+        async def execute_tool(**kwargs) -> str:
+            response = await tool_flow(**request_model(**kwargs).model_dump())
             return response.answer
 
+        # tool_flow.tool_call.name
         tool = FunctionTool(name=tool_flow.name,  # noqa
                             description=tool_flow.tool_call.description,  # noqa
-                            fn=execute_flow_async,
+                            fn=execute_tool,
                             parameters=tool_flow.tool_call.input_schema)
         self.mcp.add_tool(tool)
 
-    def __call__(self):
-        self.integrate_tool_flows()
-
+    def execute(self):
         if self.mcp_config.transport == "sse":
-            self.mcp.run(transport="sse", host=self.mcp_config.host, port=self.mcp_config.port)
-        if self.mcp_config.transport == "http":
-            self.mcp.run(transport="http", host=self.mcp_config.host, port=self.mcp_config.port)
+            self.mcp.run(transport="sse", host=self.mcp_config.host, port=self.mcp_config.port, show_banner=False)
+        elif self.mcp_config.transport == "http":
+            self.mcp.run(transport="http", host=self.mcp_config.host, port=self.mcp_config.port, show_banner=False)
         elif self.mcp_config.transport == "stdio":
-            self.mcp.run(transport="stdio")
+            self.mcp.run(transport="stdio", show_banner=False)
         else:
             raise ValueError(f"unsupported mcp transport: {self.mcp_config.transport}")

flowllm/storage/vector_store/__init__.py

@@ -1,3 +1,4 @@
 from .chroma_vector_store import ChromaVectorStore
 from .es_vector_store import EsVectorStore
 from .local_vector_store import LocalVectorStore
+from .memory_vector_store import MemoryVectorStore

flowllm/storage/vector_store/base_vector_store.py

@@ -1,9 +1,12 @@
-from abc import ABC
+import asyncio
+from abc import ABC, abstractmethod
+from functools import partial
 from pathlib import Path
-from typing import List, Iterable
+from typing import List, Iterable, Dict, Any, Optional
 
 from pydantic import BaseModel, Field
 
+from flowllm.context.service_context import C
 from flowllm.embedding_model.base_embedding_model import BaseEmbeddingModel
 from flowllm.schema.vector_node import VectorNode
 
@@ -12,33 +15,117 @@ class BaseVectorStore(BaseModel, ABC):
     embedding_model: BaseEmbeddingModel | None = Field(default=None)
     batch_size: int = Field(default=1024)
 
+    @abstractmethod
     def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        """Check if a workspace exists in the vector store."""
         raise NotImplementedError
 
-    def delete_workspace(self, workspace_id: str, **kwargs):
+    @abstractmethod
+    def delete_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Delete a workspace from the vector store."""
         raise NotImplementedError
 
-    def create_workspace(self, workspace_id: str, **kwargs):
+    @abstractmethod
+    def create_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Create a new workspace in the vector store."""
         raise NotImplementedError
 
-    def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
+    @abstractmethod
+    def iter_workspace_nodes(self, workspace_id: str, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
+        """Iterate over all nodes in a workspace."""
         raise NotImplementedError
 
-    def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
+    @abstractmethod
+    def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs) -> None:
+        """Dump workspace data to a file or path."""
         raise NotImplementedError
 
-    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, callback_fn=None,
-                       **kwargs):
+    @abstractmethod
+    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: Optional[List[VectorNode]] = None,
+                       callback_fn=None, **kwargs) -> None:
+        """Load workspace data from a file or path, or from provided nodes."""
         raise NotImplementedError
 
-    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
+    @abstractmethod
+    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs) -> None:
+        """Copy one workspace to another."""
         raise NotImplementedError
 
-    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+    @abstractmethod
+    def search(self, query: str, workspace_id: str, top_k: int = 1, filter_dict: Optional[Dict[str, Any]] = None,
+               **kwargs) -> List[VectorNode]:
+        """Search for similar vectors in the workspace."""
         raise NotImplementedError
 
-    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+    @abstractmethod
+    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs) -> None:
+        """Insert nodes into the workspace."""
         raise NotImplementedError
 
-    def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+    @abstractmethod
+    def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs) -> None:
+        """Delete nodes from the workspace by their IDs."""
         raise NotImplementedError
+
+    def close(self) -> None:
+        """Close the vector store and clean up resources. Default implementation does nothing."""
+        pass
+
+    """
+    Async versions of all methods
+    """
+
+    async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        """Async version of exist_workspace."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.exist_workspace, workspace_id, **kwargs))
+
+    async def async_delete_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Async version of delete_workspace."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.delete_workspace, workspace_id, **kwargs))
+
+    async def async_create_workspace(self, workspace_id: str, **kwargs) -> None:
+        """Async version of create_workspace."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.create_workspace, workspace_id, **kwargs))
+
+    async def async_iter_workspace_nodes(self, workspace_id: str, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
+        """Async version of iter_workspace_nodes. Returns an iterable, not an async iterator."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.iter_workspace_nodes, workspace_id,
+                                                                 callback_fn, **kwargs))
+
+    async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.dump_workspace, workspace_id, path,
+                                                                 callback_fn, **kwargs))
+
+    async def async_load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None,
+                                   callback_fn=None, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.load_workspace, workspace_id, path, nodes,
+                                                                 callback_fn, **kwargs))
+
+    async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.copy_workspace, src_workspace_id,
+                                                                 dest_workspace_id, **kwargs))
+
+    async def async_search(self, query: str, workspace_id: str, top_k: int = 1, filter_dict: dict = None,
+                           **kwargs) -> List[VectorNode]:
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.search, query, workspace_id, top_k,
+                                                                 filter_dict, **kwargs))
+
+    async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.insert, nodes, workspace_id, **kwargs))
+
+    async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, partial(self.delete, node_ids, workspace_id, **kwargs))
+
+    async def async_close(self):
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(C.thread_pool, self.close)