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

flowllm/op/search/tavily_search_op.py

@@ -0,0 +1,102 @@
+import json
+import os
+import time
+from typing import Literal
+
+from loguru import logger
+from tavily import TavilyClient
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.op.base_op import BaseOp
+from flowllm.storage.cache.data_cache import DataCache
+
+
+@C.register_op()
+class TavilySearchOp(BaseOp):
+    def __init__(self,
+                 enable_print: bool = True,
+                 enable_cache: bool = False,
+                 cache_path: str = "./web_search_cache",
+                 cache_expire_hours: float = 0.1,
+                 topic: Literal["general", "news", "finance"] = "general",
+                 max_retries: int = 3,
+                 return_only_content: bool = True,
+                 **kwargs):
+        super().__init__(**kwargs)
+
+        self.enable_print = enable_print
+        self.enable_cache = enable_cache
+        self.cache_expire_hours = cache_expire_hours
+        self.topic = topic
+        self.max_retries = max_retries
+        self.return_only_content = return_only_content
+
+        # Initialize DataCache if caching is enabled
+        self.cache = DataCache(cache_path) if self.enable_cache else None
+        self._client = TavilyClient(api_key=os.getenv("FLOW_TAVILY_API_KEY", ""))
+
+    def post_process(self, response):
+        if self.enable_print:
+            logger.info("response=\n" + json.dumps(response, indent=2, ensure_ascii=False))
+
+        return response
+
+    def execute(self):
+        # Get query from context
+        query: str = self.context.query
+
+        # Check cache first
+        if self.enable_cache and self.cache:
+            cached_result = self.cache.load(query)
+            if cached_result:
+                final_result = self.post_process(cached_result)
+                if self.return_only_content:
+                    self.context.tavily_search_result = json.dumps(final_result, ensure_ascii=False, indent=2)
+                else:
+                    self.context.tavily_search_result = final_result
+                return
+
+        for i in range(self.max_retries):
+            try:
+                response = self._client.search(query=query, topic=self.topic)
+                url_info_dict = {item["url"]: item for item in response["results"]}
+                response_extract = self._client.extract(urls=[item["url"] for item in response["results"]],
+                                                        format="text")
+
+                final_result = {}
+                for item in response_extract["results"]:
+                    url = item["url"]
+                    final_result[url] = url_info_dict[url]
+                    final_result[url]["raw_content"] = item["raw_content"]
+
+                # Cache the result if enabled
+                if self.enable_cache and self.cache:
+                    self.cache.save(query, final_result, expire_hours=self.cache_expire_hours)
+
+                final_result = self.post_process(final_result)
+
+                if self.return_only_content:
+                    self.context.tavily_search_result = json.dumps(final_result, ensure_ascii=False, indent=2)
+                else:
+                    self.context.tavily_search_result = final_result
+                return
+
+            except Exception as e:
+                logger.exception(f"tavily search with query={query} encounter error with e={e.args}")
+                time.sleep(i + 1)
+
+        self.context.tavily_search_result = "tavily search failed!"
+
+
+if __name__ == "__main__":
+    from flowllm.utils.common_utils import load_env
+
+    load_env()
+
+    C.set_default_service_config().init_by_service_config()
+
+    op = TavilySearchOp(enable_cache=True)
+    context = FlowContext(query="A股医药为什么一直涨")
+    op(context=context)
+    print(context.tavily_search_result)

flowllm/op/sequential_op.py

@@ -0,0 +1,21 @@
+from typing import List
+
+from flowllm.op.base_op import BaseOp
+
+
+class SequentialOp(BaseOp):
+
+    def __init__(self, ops: List[BaseOp], **kwargs):
+        super().__init__(**kwargs)
+        self.ops = ops
+
+    def execute(self):
+        for op in self.ops:
+            op.__call__(self.context)
+
+    def __rshift__(self, op: BaseOp):
+        if isinstance(op, SequentialOp):
+            self.ops.extend(op.ops)
+        else:
+            self.ops.append(op)
+        return self

flowllm/schema/flow_request.py

@@ -0,0 +1,12 @@
+from typing import List
+
+from pydantic import Field, BaseModel
+
+from flowllm.schema.message import Message
+
+
+class FlowRequest(BaseModel, extra="allow"):
+    query: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)
+    workspace_id: str = Field(default="")
+    metadata: dict = Field(default_factory=dict)

flowllm/schema/flow_response.py

@@ -0,0 +1,12 @@
+from typing import List
+
+from pydantic import Field, BaseModel
+
+from flowllm.schema.message import Message
+
+
+class FlowResponse(BaseModel):
+    answer: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)
+    success: bool = Field(default=True)
+    metadata: dict = Field(default_factory=dict)

flowllm/schema/message.py

@@ -0,0 +1,35 @@
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from flowllm.enumeration.role import Role
+from flowllm.schema.tool_call import ToolCall
+
+
+class Message(BaseModel):
+    role: Role = Field(default=Role.USER)
+    content: str | bytes = Field(default="")
+    reasoning_content: str = Field(default="")
+    tool_calls: List[ToolCall] = Field(default_factory=list)
+    tool_call_id: str = Field(default="")
+    metadata: dict = Field(default_factory=dict)
+
+    def simple_dump(self, add_reason_content: bool = True) -> dict:
+        result: dict
+        if self.content:
+            result = {"role": self.role.value, "content": self.content}
+        elif add_reason_content and self.reasoning_content:
+            result = {"role": self.role.value, "content": self.reasoning_content}
+        else:
+            result = {"role": self.role.value, "content": ""}
+
+        if self.tool_calls:
+            result["tool_calls"] = [x.simple_output_dump() for x in self.tool_calls]
+        return result
+
+
+class Trajectory(BaseModel):
+    task_id: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)
+    score: float = Field(default=0.0)
+    metadata: dict = Field(default_factory=dict)

flowllm/schema/service_config.py

@@ -0,0 +1,72 @@
+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=3600)
+    limit_concurrency: int = Field(default=64)
+
+
+class CmdConfig(BaseModel):
+    flow: str = Field(default="")
+    params: dict = Field(default_factory=dict)
+
+
+class FlowConfig(ToolCall):
+    flow_content: str = Field(default="")
+    service_type: str = Field(default="all", description="all/http/mcp/cmd")
+
+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="")
+    embedding_model: str = Field(default="")
+    vector_store: str = Field(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=1)
+
+    cmd: CmdConfig = Field(default_factory=CmdConfig)
+    mcp: MCPConfig = Field(default_factory=MCPConfig)
+    http: HttpConfig = Field(default_factory=HttpConfig)
+    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,118 @@
+import json
+from typing import Dict, List
+
+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")
+    enum: List[str] | None = Field(default=None, description="tool parameter enum")
+
+    def simple_dump(self) -> dict:
+        result: dict = {
+            "type": self.type,
+            "description": self.description,
+        }
+
+        if self.enum:
+            result["enum"] = self.enum
+
+        return result
+
+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: tool_param.simple_dump() 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

{llmflow → flowllm}/schema/vector_node.py

@@ -1,5 +1,6 @@
 from typing import List
 from uuid import uuid4
+
 from pydantic import BaseModel, Field
 
 

flowllm/service/__init__.py

@@ -0,0 +1,3 @@
+from .cmd_service import CmdService
+from .http_service import HttpService
+from .mcp_service import MCPService

flowllm/service/base_service.py

@@ -0,0 +1,68 @@
+from abc import abstractmethod, ABC
+from typing import Dict, Optional
+
+from loguru import logger
+from pydantic import create_model, Field
+
+from flowllm.config.pydantic_config_parser import PydanticConfigParser
+from flowllm.context.service_context import C
+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
+
+
+class BaseService(ABC):
+    TYPE_MAPPING = {
+        "str": str,
+        "int": int,
+        "float": float,
+        "bool": bool,
+        "list": list,
+        "dict": dict,
+    }
+
+    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
+        C.init_by_service_config(self.service_config)
+
+    @classmethod
+    def get_service(cls, *args, parser: type[PydanticConfigParser] = PydanticConfigParser) -> "BaseService":
+        config_parser = parser(ServiceConfig)
+        service_config: ServiceConfig = config_parser.parse_args(*args)
+        service_cls = C.resolve_service(service_config.backend)
+        return service_cls(service_config)
+
+    def _create_pydantic_model(self, flow_name: str, input_schema: Dict[str, ParamAttrs]):
+        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", __base__=FlowRequest, **fields)
+
+    def integrate_tool_flow(self, tool_flow_name: str):
+        ...
+
+    def integrate_tool_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}")
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        ...
+
+    @abstractmethod
+    def __call__(self):
+        ...

flowllm/service/cmd_service.py

@@ -0,0 +1,15 @@
+from loguru import logger
+
+from flowllm.context.service_context import C
+from flowllm.flow.gallery import CmdFlow
+from flowllm.service.base_service import BaseService
+
+
+@C.register_service("cmd")
+class CmdService(BaseService):
+
+    def __call__(self):
+        flow = CmdFlow(flow=self.service_config.cmd.flow)
+        response = flow.__call__(**self.service_config.cmd.params)
+        if response.answer:
+            logger.info(f"final_answer={response.answer}")

flowllm/service/http_service.py

@@ -0,0 +1,79 @@
+import asyncio
+from functools import partial
+
+import uvicorn
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+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.service.base_service import BaseService
+
+
+@C.register_service("http")
+class HttpService(BaseService):
+
+    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 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)
+
+        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 response
+
+        endpoint_path = f"/{tool_flow.name}"
+        self.app.post(endpoint_path, response_model=FlowResponse)(execute_endpoint)
+
+    def integrate_tool_flows(self):
+        super().integrate_tool_flows()
+
+        async def execute_endpoint() -> list:
+            loop = asyncio.get_event_loop()
+
+            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
+
+            return await loop.run_in_executor(executor=C.thread_pool, func=list_tool_flows)  # noqa
+
+        endpoint_path = "/list"
+        self.app.get(endpoint_path, response_model=list)(execute_endpoint)
+        logger.info(f"integrate endpoint={endpoint_path}")
+
+    def __call__(self):
+        self.integrate_tool_flows()
+
+        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,47 @@
+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
+
+
+@C.register_service("mcp")
+class MCPService(BaseService):
+
+    def __init__(self, *args, **kwargs):
+        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)
+
+        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
+            return response.answer
+
+        tool = FunctionTool(name=tool_flow.name,  # noqa
+                            description=tool_flow.tool_call.description,  # noqa
+                            fn=execute_flow_async,
+                            parameters=tool_flow.tool_call.input_schema)
+        self.mcp.add_tool(tool)
+
+    def __call__(self):
+        self.integrate_tool_flows()
+
+        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)
+        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/cache/__init__.py

@@ -0,0 +1 @@
+from .data_cache import DataCache