flowllm 0.1.3__py3-none-any.whl → 0.1.6__py3-none-any.whl

flowllm/op/search/tavily_search_op.py

@@ -1,109 +1,91 @@
+import asyncio
 import json
 import os
-import time
-from typing import Literal
+from functools import partial
+from typing import Union, List
 
-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
+from flowllm.context import FlowContext, C
+from flowllm.op.base_tool_op import BaseToolOp
+from flowllm.schema.tool_call import ToolCall
 
 
 @C.register_op()
-class TavilySearchOp(BaseOp):
-    def __init__(self,
-                 enable_print: bool = True,
-                 enable_cache: bool = True,
-                 cache_path: str = "./tavily_search_cache",
-                 cache_expire_hours: float = 0.1,
-                 topic: Literal["general", "news", "finance"] = "general",
-                 max_retries: int = 3,
-                 return_only_content: bool = True,
-                 **kwargs):
+class TavilySearchOp(BaseToolOp):
+    def __init__(self, save_answer: bool = False, **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._client = TavilyClient(api_key=os.getenv("FLOW_TAVILY_API_KEY", ""))
-        self.cache_path: str = cache_path
-        self._cache: DataCache | None = None
+        self.save_answer = save_answer
+        self._client: TavilyClient | None = None
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "web_search",
+            "description": "Use search keywords to retrieve relevant information from the internet. If there are multiple search keywords, please use each keyword separately to call this tool.",
+            "input_schema": {
+                "query": {
+                    "type": "str",
+                    "description": "search keyword",
+                    "required": True
+                }
+            }
+        })
 
     @property
-    def cache(self):
-        if self.enable_cache and self._cache is None:
-            self._cache = DataCache(self.cache_path)
-        return self._cache
+    def client(self):
+        if self._client is None:
+            self._client = TavilyClient(api_key=os.environ["FLOW_TAVILY_API_KEY"])
+        return self._client
 
-    def post_process(self, response):
-        if self.enable_print:
-            logger.info("response=\n" + json.dumps(response, indent=2, ensure_ascii=False))
+    def default_execute(self):
+        self.output_dict["tavily_search_result"] = "tavily search failed!"
 
-        return response
+    async def search(self, query: str):
+        loop = asyncio.get_event_loop()
+        func = partial(self.client.search, query=query)
+        task = loop.run_in_executor(executor=C.thread_pool, func=func)  # noqa
+        return await task
 
-    def execute(self):
-        # Get query from context
-        query: str = self.context.query
+    async def extract(self, urls: Union[List[str], str]):
+        loop = asyncio.get_event_loop()
+        func = partial(self.client.extract, urls=urls, format="text")
+        task = loop.run_in_executor(executor=C.thread_pool, func=func)  # noqa
+        return await task
+
+    async def async_execute(self):
+        query: str = self.input_dict["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
+                self.output_dict["tavily_search_result"] = json.dumps(cached_result, ensure_ascii=False, indent=2)
                 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
+        response = await self.search(query=query)
+        url_info_dict = {item["url"]: item for item in response["results"]}
+        response_extract = await self.extract(urls=[item["url"] for item in response["results"]])
 
-            except Exception as e:
-                logger.exception(f"tavily search with query={query} encounter error with e={e.args}")
-                time.sleep(i + 1)
+        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"]
 
-        self.context.tavily_search_result = "tavily search failed!"
+        if self.enable_cache and self.cache is not None:
+            self.cache.save(query, final_result, expire_hours=self.cache_expire_hours)
 
+        self.output_dict["tavily_search_result"] = json.dumps(final_result, ensure_ascii=False, indent=2)
+        if self.save_answer:
+            self.context.response.answer = self.output_dict["tavily_search_result"]
 
-if __name__ == "__main__":
-    from flowllm.utils.common_utils import load_env
+async def async_main():
+    C.set_service_config().init_by_service_config()
 
-    load_env()
+    op = TavilySearchOp()
+    context = FlowContext(query="what is AI?")
+    await op.async_call(context=context)
+    print(context.tavily_search_result)
 
-    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)
+if __name__ == "__main__":
+    asyncio.run(async_main())

flowllm/op/sequential_op.py

@@ -13,6 +13,10 @@ class SequentialOp(BaseOp):
         for op in self.ops:
             op.__call__(self.context)
 
+    async def async_execute(self):
+        for op in self.ops:
+            await op.async_call(self.context)
+
     def __rshift__(self, op: BaseOp):
         if isinstance(op, SequentialOp):
             self.ops.extend(op.ops)

flowllm/schema/flow_stream_chunk.py

@@ -0,0 +1,11 @@
+from pydantic import Field, BaseModel
+
+from flowllm.enumeration.chunk_enum import ChunkEnum
+
+
+class FlowStreamChunk(BaseModel):
+    flow_id: str = Field(default="")
+    chunk_type: ChunkEnum = Field(default=ChunkEnum.ANSWER)
+    chunk: str | bytes = Field(default="")
+    done: bool = Field(default=False)
+    metadata: dict = Field(default_factory=dict)

flowllm/schema/service_config.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, List
 
 from pydantic import BaseModel, Field
 
@@ -6,7 +6,7 @@ from flowllm.schema.tool_call import ToolCall
 
 
 class MCPConfig(BaseModel):
-    transport: str = Field(default="", description="stdio/http/sse/streamable-http")
+    transport: str = Field(default="", description="stdio/http/sse")
     host: str = Field(default="0.0.0.0")
     port: int = Field(default=8001)
 
@@ -25,11 +25,14 @@ class CmdConfig(BaseModel):
 
 class FlowConfig(ToolCall):
     flow_content: str = Field(default="")
-    service_type: str = Field(default="all", description="all/http/mcp/cmd")
+    use_async: bool = Field(default=True)
+    service_type: str = Field(default="http+mcp", description="http+mcp/http/mcp")
+    stream: bool = Field(default=False)
 
 class OpConfig(BaseModel):
     backend: str = Field(default="")
     language: str = Field(default="")
+    max_retries: int = Field(default=1)
     raise_exception: bool = Field(default=True)
     prompt_path: str = Field(default="")
     llm: str = Field(default="")
@@ -61,6 +64,8 @@ class ServiceConfig(BaseModel):
     language: str = Field(default="")
     thread_pool_max_workers: int = Field(default=16)
     ray_max_workers: int = Field(default=1)
+    import_config: str = Field(default="", description="Import the configuration in the same path as the base")
+    disabled_flows: List[str] = Field(default_factory=list)
 
     cmd: CmdConfig = Field(default_factory=CmdConfig)
     mcp: MCPConfig = Field(default_factory=MCPConfig)

flowllm/schema/tool_call.py

@@ -1,7 +1,8 @@
 import json
 from typing import Dict, List
 
-from pydantic import BaseModel, Field
+from mcp.types import Tool
+from pydantic import BaseModel, Field, model_validator
 
 
 class ParamAttrs(BaseModel):
@@ -64,6 +65,17 @@ class ToolCall(BaseModel):
     input_schema: Dict[str, ParamAttrs] = Field(default_factory=dict)
     output_schema: Dict[str, ParamAttrs] = Field(default_factory=dict)
 
+    @model_validator(mode="before")
+    @classmethod
+    def init_tool_call(cls, data: dict):
+        tool_type = data.get("type", "")
+        tool_type_dict = data.get(tool_type, {})
+
+        for key in ["name", "arguments"]:
+            if key not in data:
+                data[key] = tool_type_dict.get(key, "")
+        return data
+
     @property
     def argument_dict(self) -> dict:
         return json.loads(self.arguments)
@@ -120,3 +132,36 @@ class ToolCall(BaseModel):
             raise NotImplementedError(f"version {version} not supported")
 
         return self
+
+    @classmethod
+    def from_mcp_tool(cls, tool: Tool) -> "ToolCall":
+        input_schema = {}
+        properties = tool.inputSchema["properties"]
+        required = tool.inputSchema["required"]
+        for name, attr_dict in properties.items():
+            param_attrs = ParamAttrs()
+
+            if name in required:
+                param_attrs.required = True
+            param_attrs.type = attr_dict.get("type", "str")
+            param_attrs.description = attr_dict.get("description", "")
+            if "enum" in attr_dict:
+                param_attrs.enum = attr_dict["enum"]
+            input_schema[name] = param_attrs
+
+        return cls(name=tool.name,
+                   description=tool.description,
+                   input_schema=input_schema)
+
+
+if __name__ == "__main__":
+    tool_call = ToolCall(**{
+        "id": "call_0fb6077ad56f4647b0b04a",
+        "function": {
+            "arguments": "{\"symbol\": \"ZETA\"}",
+            "name": "get_stock_info"
+        },
+        "type": "function",
+        "index": 0
+    })
+    print(tool_call.simple_output_dump())

flowllm/service/__init__.py

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

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)
+
+        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
 
         endpoint_path = "list"
-        self.app.get("/" + endpoint_path, response_model=list)(execute_endpoint)
+        self.app.get(f"/{endpoint_path}", response_model=list)(execute_endpoint)
         logger.info(f"integrate endpoint={endpoint_path}")
 
-    def __call__(self):
-        self.integrate_tool_flows()
-
+    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