flowllm 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

flowllm/schema/service_config.py

@@ -0,0 +1,76 @@
+from typing import Dict
+
+from pydantic import BaseModel, Field
+
+from flowllm.schema.tool_call import ToolCall
+
+
+class MCPConfig(BaseModel):
+    transport: str = Field(default="", description="stdio/http/sse/streamable-http")
+    host: str = Field(default="0.0.0.0")
+    port: int = Field(default=8001)
+
+
+class HttpConfig(BaseModel):
+    host: str = Field(default="0.0.0.0")
+    port: int = Field(default=8001)
+    timeout_keep_alive: int = Field(default=600)
+    limit_concurrency: int = Field(default=64)
+
+
+class FlowConfig(ToolCall):
+    flow_content: str = Field(default="")
+
+    def set_name(self, name: str):
+        self.name = name
+        return self
+
+
+class FlowEngineConfig(BaseModel):
+    backend: str = Field(default="")
+    params: dict = Field(default_factory=dict)
+
+
+class OpConfig(BaseModel):
+    backend: str = Field(default="")
+    language: str = Field(default="")
+    raise_exception: bool = Field(default=True)
+    prompt_path: str = Field(default="")
+    llm: str = Field(default="default")
+    embedding_model: str = Field(default="default")
+    vector_store: str = Field(default="default")
+    params: dict = Field(default_factory=dict)
+
+
+class LLMConfig(BaseModel):
+    backend: str = Field(default="")
+    model_name: str = Field(default="")
+    params: dict = Field(default_factory=dict)
+
+
+class EmbeddingModelConfig(BaseModel):
+    backend: str = Field(default="")
+    model_name: str = Field(default="")
+    params: dict = Field(default_factory=dict)
+
+
+class VectorStoreConfig(BaseModel):
+    backend: str = Field(default="")
+    embedding_model: str = Field(default="")
+    params: dict = Field(default_factory=dict)
+
+
+class ServiceConfig(BaseModel):
+    backend: str = Field(default="")
+    language: str = Field(default="")
+    thread_pool_max_workers: int = Field(default=16)
+    ray_max_workers: int = Field(default=8)
+
+    mcp: MCPConfig = Field(default_factory=MCPConfig)
+    http: HttpConfig = Field(default_factory=HttpConfig)
+    flow_engine: FlowEngineConfig = Field(default_factory=FlowEngineConfig)
+    flow: Dict[str, FlowConfig] = Field(default_factory=dict)
+    op: Dict[str, OpConfig] = Field(default_factory=dict)
+    llm: Dict[str, LLMConfig] = Field(default_factory=dict)
+    embedding_model: Dict[str, EmbeddingModelConfig] = Field(default_factory=dict)
+    vector_store: Dict[str, VectorStoreConfig] = Field(default_factory=dict)

flowllm/schema/tool_call.py

@@ -0,0 +1,110 @@
+import json
+from typing import Dict
+
+from pydantic import BaseModel, Field
+
+
+class ParamAttrs(BaseModel):
+    type: str = Field(default="str", description="tool parameter type")
+    description: str = Field(default="", description="tool parameter description")
+    required: bool = Field(default=True, description="tool parameter required")
+
+
+class ToolCall(BaseModel):
+    """
+    input:
+    {
+        "type": "function",
+        "function": {
+            "name": "get_current_weather",
+            "description": "It is very useful when you want to check the weather of a specified city.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.",
+                    }
+                },
+                "required": ["location"]
+            }
+        }
+    }
+    output:
+    {
+        "index": 0
+        "id": "call_6596dafa2a6a46f7a217da",
+        "function": {
+            "arguments": "{\"location\": \"Beijing\"}",
+            "name": "get_current_weather"
+        },
+        "type": "function",
+    }
+    """
+
+    index: int = Field(default=0)
+    id: str = Field(default="")
+    type: str = Field(default="function")
+    name: str = Field(default="")
+
+    arguments: dict = Field(default_factory=dict, description="tool execution arguments")
+
+    description: str = Field(default="")
+    input_schema: Dict[str, ParamAttrs] = Field(default_factory=dict)
+    output_schema: Dict[str, ParamAttrs] = Field(default_factory=dict)
+
+    def simple_input_dump(self, version: str = "v1") -> dict:
+        if version == "v1":
+            required_list = [name for name, tool_param in self.input_schema.items() if tool_param.required]
+            properties = {name: {
+                "type": tool_param.type,
+                "description": tool_param.description
+            } for name, tool_param in self.input_schema.items()}
+
+            return {
+                "type": self.type,
+                self.type: {
+                    "name": self.name,
+                    "description": self.description,
+                    "parameters": {
+                        "type": "object",
+                        "properties": properties,
+                        "required": required_list
+                    },
+                },
+            }
+
+        else:
+            raise NotImplementedError(f"version {version} not supported")
+
+    def simple_output_dump(self, version: str = "v1") -> dict:
+        if version == "v1":
+            return {
+                "index": self.index,
+                "id": self.id,
+                self.type: {
+                    "arguments": json.dumps(self.arguments, ensure_ascii=False),
+                    "name": self.name
+                },
+                "type": self.type,
+            }
+        else:
+            raise NotImplementedError(f"version {version} not supported")
+
+    def update_by_output(self, data: dict, version: str = "v1"):
+        if version == "v1":
+            self.index = data.get("index", 0)
+            self.id = data.get("id", "")
+            tool_type = data.get("type", "")
+            tool_type_dict = data.get(tool_type, {})
+            if tool_type_dict:
+                name = tool_type_dict.get("name", "")
+                arguments = tool_type_dict.get("arguments", "")
+                if name:
+                    self.name = name
+                if arguments:
+                    self.arguments = json.loads(arguments)
+        else:
+            raise NotImplementedError(f"version {version} not supported")
+
+        return self

flowllm/service/__init__.py

@@ -0,0 +1,2 @@
+from flowllm.service.mcp_service import MCPService
+from flowllm.service.http_service import HttpService

flowllm/service/base_service.py

@@ -0,0 +1,59 @@
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict
+
+from loguru import logger
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.flow_engine.base_flow_engine import BaseFlowEngine
+from flowllm.schema.flow_response import FlowResponse
+from flowllm.schema.service_config import ServiceConfig, EmbeddingModelConfig, FlowConfig
+
+
+class BaseService:
+
+    def __init__(self, service_config: ServiceConfig):
+        self.service_config = service_config
+
+        C.language = self.service_config.language
+        C.thread_pool = ThreadPoolExecutor(max_workers=self.service_config.thread_pool_max_workers)
+        for name, config in self.service_config.vector_store.items():
+            vector_store_cls = C.resolve_vector_store(config.backend)
+            embedding_model_config: EmbeddingModelConfig = self.service_config.embedding_model[config.embedding_model]
+            embedding_model_cls = C.resolve_embedding_model(embedding_model_config.backend)
+            embedding_model = embedding_model_cls(model_name=embedding_model_config.model_name,
+                                                  **embedding_model_config.params)
+            C.set_vector_store(name, vector_store_cls(embedding_model=embedding_model, **config.params))
+
+        self.flow_engine_config = self.service_config.flow_engine
+        self.flow_engine_cls = C.resolve_flow_engine(self.flow_engine_config.backend)
+        self.flow_config_dict: Dict[str, FlowConfig] = \
+            {name: config.set_name(name) for name, config in self.service_config.flow.items()}
+
+        self.mcp_config = self.service_config.mcp
+        self.http_config = self.service_config.http
+
+    def execute_flow(self, flow_name: str, **kwargs) -> FlowResponse:
+        response = FlowResponse()
+        try:
+            logger.info(f"request.params={kwargs}")
+            flow_context = FlowContext(**kwargs,
+                                       response=response,
+                                       service_config=self.service_config.model_copy(deep=True))
+
+            flow_config = self.flow_config_dict[flow_name]
+            flow_engine: BaseFlowEngine = self.flow_engine_cls(flow_name=flow_name,
+                                                               flow_content=flow_config.flow_content,
+                                                               flow_context=flow_context,
+                                                               **self.flow_engine_config.params)
+            flow_engine()
+
+        except Exception as e:
+            logger.exception(f"flow_name={flow_name} encounter error={e.args}")
+            response.success = False
+            response.answer = str(e.args)
+
+        return response
+
+    def __call__(self):
+        raise NotImplementedError

flowllm/service/http_service.py

@@ -0,0 +1,87 @@
+import asyncio
+from functools import partial
+from typing import Dict, Optional
+
+import uvicorn
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from loguru import logger
+from pydantic import BaseModel, create_model, Field
+
+from flowllm.context.service_context import C
+from flowllm.schema.flow_response import FlowResponse
+from flowllm.schema.tool_call import ParamAttrs
+from flowllm.service.base_service import BaseService
+from flowllm.utils.common_utils import snake_to_camel
+
+
+@C.register_service("http")
+class HttpService(BaseService):
+    TYPE_MAPPING = {
+        "str": str,
+        "int": int,
+        "float": float,
+        "bool": bool
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.app = FastAPI(title="FlowLLM", description="HTTP API for FlowLLM")
+        
+        # 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 _create_pydantic_model(self, flow_name: str, input_schema: Dict[str, ParamAttrs]) -> BaseModel:
+        # Create a dynamic Pydantic model based on flow input schema
+        fields = {}
+
+        for param_name, param_config in input_schema.items():
+            field_type = self.TYPE_MAPPING.get(param_config.type, str)
+
+            if not param_config.required:
+                fields[param_name] = (Optional[field_type], Field(default=None, description=param_config.description))
+            else:
+                fields[param_name] = (field_type, Field(default=..., description=param_config.description))
+
+        return create_model(f"{snake_to_camel(flow_name)}Model", **fields)
+    
+    def register_flow(self, flow_name: str):
+        """Register a flow as an HTTP endpoint"""
+        flow_config = self.flow_config_dict[flow_name]
+        request_model = self._create_pydantic_model(flow_name, flow_config.input_schema)
+
+        async def execute_flow_endpoint(request: request_model) -> FlowResponse:
+            loop = asyncio.get_event_loop()
+            response: FlowResponse = await loop.run_in_executor(
+                executor=C.thread_pool,
+                func=partial(self.execute_flow, flow_name=flow_name, **request.model_dump()))  # noqa
+
+            return response
+
+        endpoint_path = f"/{flow_name}"
+        self.app.post(endpoint_path, response_model=FlowResponse)(execute_flow_endpoint)
+        logger.info(f"register flow={flow_name} endpoint={endpoint_path}")
+    
+    def __call__(self):
+        for flow_name in self.flow_config_dict:
+            self.register_flow(flow_name)
+        
+        # Start the server
+        uvicorn.run(self.app,
+                    host=self.http_config.host,
+                    port=self.http_config.port,
+                    timeout_keep_alive=self.http_config.timeout_keep_alive,
+                    limit_concurrency=self.http_config.limit_concurrency)

flowllm/service/mcp_service.py

@@ -0,0 +1,45 @@
+import asyncio
+from functools import partial
+
+from fastmcp import FastMCP
+from fastmcp.tools import FunctionTool
+from loguru import logger
+
+from flowllm.context.service_context import C
+from flowllm.service.base_service import BaseService
+
+
+@C.register_service("mcp")
+class MCPService(BaseService):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mcp = FastMCP("FlowLLM")
+
+    def register_flow(self, flow_name: str):
+        flow_config = self.flow_config_dict[flow_name]
+
+        async def execute_flow_async(**kwargs) -> str:
+            loop = asyncio.get_event_loop()
+            response = await loop.run_in_executor(
+                executor=C.thread_pool,
+                func=partial(self.execute_flow, flow_name=flow_name, **kwargs))  # noqa
+            return response.answer
+
+        tool = FunctionTool(name=flow_name,  # noqa
+                            description=flow_config.description,  # noqa
+                            fn=execute_flow_async,
+                            parameters=flow_config.input_schema)
+        self.mcp.add_tool(tool)
+        logger.info(f"register flow={flow_name}")
+
+    def __call__(self):
+        for flow_name in self.flow_config_dict:
+            self.register_flow(flow_name)
+
+        if self.mcp_config.transport == "sse":
+            self.mcp.run(transport="sse", host=self.mcp_config.host, port=self.mcp_config.port)
+        elif self.mcp_config.transport == "stdio":
+            self.mcp.run(transport="stdio")
+        else:
+            raise ValueError(f"unsupported mcp transport: {self.mcp_config.transport}")

flowllm/storage/__init__.py

@@ -0,0 +1 @@
+from flowllm.storage import vector_store

flowllm/storage/vector_store/__init__.py

@@ -0,0 +1,3 @@
+from flowllm.storage.vector_store.chroma_vector_store import ChromaVectorStore
+from flowllm.storage.vector_store.es_vector_store import EsVectorStore
+from flowllm.storage.vector_store.local_vector_store import LocalVectorStore

flowllm/storage/vector_store/base_vector_store.py

@@ -0,0 +1,44 @@
+from abc import ABC
+from pathlib import Path
+from typing import List, Iterable
+
+from pydantic import BaseModel, Field
+
+from flowllm.embedding_model.base_embedding_model import BaseEmbeddingModel
+from flowllm.schema.vector_node import VectorNode
+
+
+class BaseVectorStore(BaseModel, ABC):
+    embedding_model: BaseEmbeddingModel | None = Field(default=None)
+    batch_size: int = Field(default=1024)
+
+    def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
+        raise NotImplementedError
+
+    def delete_workspace(self, workspace_id: str, **kwargs):
+        raise NotImplementedError
+
+    def create_workspace(self, workspace_id: str, **kwargs):
+        raise NotImplementedError
+
+    def _iter_workspace_nodes(self, workspace_id: str, **kwargs) -> Iterable[VectorNode]:
+        raise NotImplementedError
+
+    def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs):
+        raise NotImplementedError
+
+    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, callback_fn=None,
+                       **kwargs):
+        raise NotImplementedError
+
+    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
+        raise NotImplementedError
+
+    def search(self, query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode]:
+        raise NotImplementedError
+
+    def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs):
+        raise NotImplementedError
+
+    def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs):
+        raise NotImplementedError

{llmflow → flowllm/storage}/vector_store/chroma_vector_store.py

@@ -6,17 +6,16 @@ from chromadb.config import Settings
 from loguru import logger
 from pydantic import Field, PrivateAttr, model_validator
 
-from llmflow.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
-from llmflow.schema.vector_node import VectorNode
-from llmflow.vector_store import VECTOR_STORE_REGISTRY
-from llmflow.vector_store.base_vector_store import BaseVectorStore
+from flowllm.context.service_context import C
+from flowllm.schema.vector_node import VectorNode
+from flowllm.storage.vector_store.local_vector_store import LocalVectorStore
 
 
-@VECTOR_STORE_REGISTRY.register("chroma")
-class ChromaVectorStore(BaseVectorStore):
+@C.register_vector_store("chroma")
+class ChromaVectorStore(LocalVectorStore):
     store_dir: str = Field(default="./chroma_vector_store")
     collections: dict = Field(default_factory=dict)
-    _client: chromadb.Client = PrivateAttr()
+    _client: chromadb.ClientAPI = PrivateAttr()
 
     @model_validator(mode="after")
     def init_client(self):
@@ -97,8 +96,10 @@ class ChromaVectorStore(BaseVectorStore):
 
 
 def main():
-    from dotenv import load_dotenv
-    load_dotenv()
+    from flowllm.utils.common_utils import load_env
+    from flowllm.embedding_model import OpenAICompatibleEmbeddingModel
+
+    load_env()
 
     embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
     workspace_id = "chroma_test_index"
@@ -185,4 +186,4 @@ def main():
 
 if __name__ == "__main__":
     main()
-    # launch with: python -m llmflow.storage.chroma_vector_store
+    # launch with: python -m flowllm.storage.chroma_vector_store

{llmflow → flowllm/storage}/vector_store/es_vector_store.py

@@ -6,14 +6,13 @@ from elasticsearch.helpers import bulk
 from loguru import logger
 from pydantic import Field, PrivateAttr, model_validator
 
-from llmflow.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
-from llmflow.schema.vector_node import VectorNode
-from llmflow.vector_store import VECTOR_STORE_REGISTRY
-from llmflow.vector_store.base_vector_store import BaseVectorStore
+from flowllm.context.service_context import C
+from flowllm.schema.vector_node import VectorNode
+from flowllm.storage.vector_store.local_vector_store import LocalVectorStore
 
 
-@VECTOR_STORE_REGISTRY.register("elasticsearch")
-class EsVectorStore(BaseVectorStore):
+@C.register_vector_store("elasticsearch")
+class EsVectorStore(LocalVectorStore):
     hosts: str | List[str] = Field(default_factory=lambda: os.getenv("ES_HOSTS", "http://localhost:9200"))
     basic_auth: str | Tuple[str, str] | None = Field(default=None)
     retrieve_filters: List[dict] = []
@@ -24,6 +23,7 @@ class EsVectorStore(BaseVectorStore):
         if isinstance(self.hosts, str):
             self.hosts = [self.hosts]
         self._client = Elasticsearch(hosts=self.hosts, basic_auth=self.basic_auth)
+        logger.info(f"Elasticsearch client initialized with hosts: {self.hosts}")
         return self
 
     def exist_workspace(self, workspace_id: str, **kwargs) -> bool:
@@ -162,8 +162,10 @@ class EsVectorStore(BaseVectorStore):
 
 
 def main():
-    from dotenv import load_dotenv
-    load_dotenv()
+    from flowllm.utils.common_utils import load_env
+    from flowllm.embedding_model import OpenAICompatibleEmbeddingModel
+
+    load_env()
 
     embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
     workspace_id = "rag_nodes_index"
@@ -224,4 +226,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-    # launch with: python -m llmflow.storage.es_vector_store