flowllm 0.1.0__py3-none-any.whl

llmflow/tool/mcp_tool.py

@@ -0,0 +1,77 @@
+import asyncio
+from typing import List
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+from pydantic import Field, model_validator
+
+from llmflow.tool import TOOL_REGISTRY
+from llmflow.tool.base_tool import BaseTool
+
+
+@TOOL_REGISTRY.register()
+class MCPTool(BaseTool):
+    server_url: str = Field(..., description="MCP server URL")
+    tool_name_list: List[str] = Field(default_factory=list)
+    cache_tools: dict = Field(default_factory=dict, alias="cache_tools")
+
+    @model_validator(mode="after")
+    def refresh_tools(self):
+        self.refresh()
+        return self
+
+    async def _get_tools(self):
+        async with sse_client(url=self.server_url) as streams:
+            async with ClientSession(streams[0], streams[1]) as session:
+                await session.initialize()
+                tools = await session.list_tools()
+        return tools
+
+    def refresh(self):
+        self.tool_name_list.clear()
+        self.cache_tools.clear()
+
+        if "sse" in self.server_url:
+            original_tool_list = asyncio.run(self._get_tools())
+            for tool in original_tool_list.tools:
+                self.cache_tools[tool.name] = tool
+                self.tool_name_list.append(tool.name)
+        else:
+            raise NotImplementedError("Non-SSE refresh not implemented yet")
+
+    @property
+    def input_schema(self) -> dict:
+        return {x: self.cache_tools[x].inputSchema for x in self.cache_tools}
+
+    @property
+    def output_schema(self) -> dict:
+        raise NotImplementedError("Output schema not implemented yet")
+
+    def get_tool_description(self, tool_name: str, schema: bool = False) -> str:
+        if tool_name not in self.cache_tools:
+            raise RuntimeError(f"Tool {tool_name} not found")
+
+        tool = self.cache_tools.get(tool_name)
+        description = f"tool={tool_name} description={tool.description}\n"
+        if schema:
+            description += f"input_schema={self.input_schema[tool_name]}\n" \
+                           f"output_schema={self.output_schema[tool_name]}\n"
+        return description.strip()
+
+    async def async_execute(self, tool_name: str, **kwargs):
+        if "sse" in self.server_url:
+            async with sse_client(url=self.server_url) as streams:
+                async with ClientSession(streams[0], streams[1]) as session:
+                    await session.initialize()
+                    results = await session.call_tool(tool_name, kwargs)
+            return results.content[0].text, results.isError
+
+        else:
+            raise NotImplementedError("Non-SSE execute not implemented yet")
+
+    def _execute(self, **kwargs):
+        return asyncio.run(self.async_execute(**kwargs))
+
+    def get_cache_id(self, **kwargs) -> str:
+        # Implement a method to generate a unique cache ID based on the input
+        return f"{kwargs.get('tool_name')}_{hash(frozenset(kwargs.get('args', {}).items()))}"

llmflow/tool/tavily_search_tool.py

@@ -0,0 +1,109 @@
+import json
+import os
+import re
+import time
+from typing import Literal
+
+from loguru import logger
+from pydantic import Field, model_validator, PrivateAttr
+from tavily import TavilyClient
+
+from llmflow.tool import TOOL_REGISTRY
+from llmflow.tool.base_tool import BaseTool
+
+
+@TOOL_REGISTRY.register()
+class TavilySearchTool(BaseTool):
+    name: str = "web_search"
+    description: str = "Use query to retrieve relevant information from the internet."
+    parameters: dict = {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "search query",
+            }
+        },
+        "required": ["query"]
+    }
+    enable_print: bool = Field(default=True)
+    enable_cache: bool = Field(default=False)
+    cache_path: str = Field(default="./web_search_cache")
+    topic: Literal["general", "news", "finance"] = Field(default="general", description="finance, general")
+
+    _client: TavilyClient | None = PrivateAttr()
+
+    @model_validator(mode="after")
+    def init(self):
+        if not os.path.exists(self.cache_path):
+            os.makedirs(self.cache_path)
+
+        self._client = TavilyClient()
+        return self
+
+    def load_cache(self, cache_name: str = "default") -> dict:
+        cache_file = os.path.join(self.cache_path, cache_name + ".jsonl")
+        if not os.path.exists(cache_file):
+            return {}
+
+        with open(cache_file) as f:
+            return json.load(f)
+
+    def dump_cache(self, cache_dict: dict, cache_name: str = "default"):
+        cache_file = os.path.join(self.cache_path, cache_name + ".jsonl")
+        with open(cache_file, "w") as f:
+            return json.dump(cache_dict, f, indent=2, ensure_ascii=False)
+
+    @staticmethod
+    def remove_urls_and_images(text):
+        pattern = re.compile(r'https?://[-A-Za-z0-9+&@#/%?=~_|!:,.;]+[-A-Za-z0-9+&@#/%=~_|]')
+        result = pattern.sub("", text)
+        return result
+
+    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, query: str = "", **kwargs):
+        assert query, "Query cannot be empty"
+
+        cache_dict = {}
+        if self.enable_cache:
+            cache_dict = self.load_cache()
+            if query in cache_dict:
+                return self.post_process(cache_dict[query])
+
+        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"]
+
+                if self.enable_cache:
+                    cache_dict[query] = final_result
+                    self.dump_cache(cache_dict)
+
+                return self.post_process(final_result)
+
+            except Exception as e:
+                logger.exception(f"tavily search with query={query} encounter error with e={e.args}")
+                time.sleep(i + 1)
+
+        return None
+
+
+if __name__ == "__main__":
+    from dotenv import load_dotenv
+
+    load_dotenv()
+    tool = TavilySearchTool()
+    tool.execute(query="A股医药为什么一直涨")

llmflow/tool/terminate_tool.py

@@ -0,0 +1,23 @@
+from llmflow.tool import TOOL_REGISTRY
+from llmflow.tool.base_tool import BaseTool
+
+
+@TOOL_REGISTRY.register()
+class TerminateTool(BaseTool):
+    name: str = "terminate"
+    description: str = "If you can answer the user's question based on the context, be sure to use the **terminate** tool."
+    parameters: dict = {
+        "type": "object",
+        "properties": {
+            "status": {
+                "type": "string",
+                "description": "Please determine whether the user's question has been completed. (success / failure)",
+                "enum": ["success", "failure"],
+            }
+        },
+        "required": ["status"],
+    }
+
+    def execute(self, status: str):
+        self.success = status in ["success", "failure"]
+        return f"The interaction has been completed with status: {status}"

llmflow/utils/common_utils.py

@@ -0,0 +1,17 @@
+import re
+
+
+def camel_to_snake(content: str) -> str:
+    """
+    BaseWorker -> base_worker
+    """
+    snake_str = re.sub(r'(?<!^)(?=[A-Z])', '_', content).lower()
+    return snake_str
+
+
+def snake_to_camel(content: str) -> str:
+    """
+    base_worker -> BaseWorker
+    """
+    camel_str = "".join(x.capitalize() for x in content.split("_"))
+    return camel_str

llmflow/utils/file_handler.py

@@ -0,0 +1,25 @@
+import json
+from pathlib import Path
+
+import yaml
+
+
+class FileHandler:
+
+    def __init__(self, file_path: str | Path):
+        self.file_path: Path = Path(file_path)
+        suffix = Path(self.file_path).suffix
+        if suffix == ".json":
+            self._obj = json
+        elif suffix == ".yaml":
+            self._obj = yaml
+        else:
+            raise ValueError(f"unsupported file type={suffix}")
+
+    def dump(self, config, **kwargs):
+        with open(self.file_path, "w") as f:
+            self._obj.dump(config, f, **kwargs)
+
+    def load(self, **kwargs):
+        with open(self.file_path, "r") as f:
+            return self._obj.load(f, **kwargs)

llmflow/utils/http_client.py

@@ -0,0 +1,156 @@
+import http
+import time
+from typing import Any
+
+import requests
+from loguru import logger
+from pydantic import BaseModel, Field, PrivateAttr, model_validator
+
+from llmflow.enumeration.http_enum import HttpEnum
+
+
+class HttpClient(BaseModel):
+    url: str = Field(default="")
+    keep_alive: bool = Field(default=False, description="if true, use session to keep long connection")
+    timeout: int = Field(default=300, description="request timeout, second")
+
+    return_default_if_error: bool = Field(default=True)
+    request_start_time: float = Field(default_factory=time.time)
+    request_time_cost: float = Field(default=0.0, description="request time cost")
+
+    retry_sleep_time: float = Field(default=0.5, description="interval time for retry")
+    retry_time_multiplier: float = Field(default=2.0, description="retry time multiplier")
+    retry_max_count: int = Field(default=1, description="maximum number of retries")
+
+    _client: Any = PrivateAttr()
+
+    @model_validator(mode="after")
+    def init_client(self):
+        self._client = requests.Session() if self.keep_alive else requests
+        return self
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+        self.request_time_cost: float = time.time() - self.request_start_time
+
+    def close(self):
+        if isinstance(self._client, requests.Session):
+            self._client.close()
+
+    def _request(self,
+                 data: str = None,
+                 json_data: dict = None,
+                 headers: dict = None,
+                 stream: bool = False,
+                 http_enum: HttpEnum | str = HttpEnum.POST):
+
+        if isinstance(http_enum, str):
+            http_enum = HttpEnum(http_enum)
+
+        if http_enum is HttpEnum.POST:
+            response: requests.Response = self._client.post(url=self.url,
+                                                            data=data,
+                                                            json=json_data,
+                                                            headers=headers,
+                                                            stream=stream,
+                                                            timeout=self.timeout)
+
+        elif http_enum is HttpEnum.GET:
+            response: requests.Response = self._client.get(url=self.url,
+                                                           data=data,
+                                                           json=json_data,
+                                                           headers=headers,
+                                                           stream=stream,
+                                                           timeout=self.timeout)
+
+        else:
+            raise NotImplementedError
+
+        if response.status_code != http.HTTPStatus.OK:
+            raise RuntimeError(f"request failed! content={response.json()}")
+
+        return response
+
+    def parse_result(self, response: requests.Response | Any = None, **kwargs):
+        return response.json()
+
+    def return_default(self, **kwargs):
+        return None
+
+    def request(self,
+                data: str | Any = None,
+                json_data: dict = None,
+                headers: dict = None,
+                http_enum: HttpEnum | str = HttpEnum.POST,
+                **kwargs):
+
+        retry_sleep_time = self.retry_sleep_time
+        for i in range(self.retry_max_count):
+            try:
+                response = self._request(data=data, json_data=json_data, headers=headers, http_enum=http_enum)
+                result = self.parse_result(response=response,
+                                           data=data,
+                                           json_data=json_data,
+                                           headers=headers,
+                                           http_enum=http_enum,
+                                           **kwargs)
+                return result
+
+            except Exception as e:
+                logger.exception(f"{self.__class__.__name__} {i}th request failed with args={e.args}")
+
+                if i == self.retry_max_count - 1:
+                    if self.return_default_if_error:
+                        return self.return_default()
+                    else:
+                        raise e
+
+                retry_sleep_time *= self.retry_time_multiplier
+                time.sleep(retry_sleep_time)
+
+        return None
+
+    def request_stream(self,
+                       data: str = None,
+                       json_data: dict = None,
+                       headers: dict = None,
+                       http_enum: HttpEnum | str = HttpEnum.POST,
+                       **kwargs):
+
+        retry_sleep_time = self.retry_sleep_time
+        for i in range(self.retry_max_count):
+            try:
+                response = self._request(data=data,
+                                         json_data=json_data,
+                                         headers=headers,
+                                         stream=True,
+                                         http_enum=http_enum)
+                request_context = {}
+                for iter_idx, line in enumerate(response.iter_lines()):
+                    yield self.parse_result(line=line,
+                                            request_context=request_context,
+                                            index=iter_idx,
+                                            data=data,
+                                            json_data=json_data,
+                                            headers=headers,
+                                            http_enum=http_enum,
+                                            **kwargs)
+
+                return None
+
+            except Exception as e:
+                logger.exception(f"{self.__class__.__name__} {i}th request failed with args={e.args}")
+
+                if i == self.retry_max_count - 1:
+                    if self.return_default_if_error:
+                        return self.return_default()
+                    else:
+                        raise e
+
+                retry_sleep_time *= self.retry_time_multiplier
+                time.sleep(retry_sleep_time)
+
+        return None

llmflow/utils/op_utils.py

@@ -0,0 +1,102 @@
+from typing import List
+
+from llmflow.enumeration.role import Role
+from llmflow.schema.message import Message, Trajectory
+import json
+import re
+from loguru import logger
+
+def merge_messages_content(messages: List[Message | dict]) -> str:
+    content_collector = []
+    for i, message in enumerate(messages):
+        if isinstance(message, dict):
+            message = Message(**message)
+
+        if message.role is Role.ASSISTANT:
+            line = f"### step.{i} role={message.role.value} content=\n{message.reasoning_content}\n\n{message.content}\n"
+            if message.tool_calls:
+                for tool_call in message.tool_calls:
+                    line += f" - tool call={tool_call.name}\n   params={tool_call.arguments}\n"
+            content_collector.append(line)
+
+        elif message.role is Role.USER:
+            line = f"### step.{i} role={message.role.value} content=\n{message.content}\n"
+            content_collector.append(line)
+
+        elif message.role is Role.TOOL:
+            line = f"### step.{i} role={message.role.value} tool call result=\n{message.content}\n"
+            content_collector.append(line)
+
+    return "\n".join(content_collector)
+
+
+def parse_json_experience_response(response: str) -> List[dict]:
+    """Parse JSON formatted experience response"""
+    try:
+        # Extract JSON blocks
+        json_pattern = r'```json\s*([\s\S]*?)\s*```'
+        json_blocks = re.findall(json_pattern, response)
+
+        if json_blocks:
+            parsed = json.loads(json_blocks[0])
+
+            # Handle array format
+            if isinstance(parsed, list):
+                experiences = []
+                for exp_data in parsed:
+                    if isinstance(exp_data, dict) and (
+                            ("when_to_use" in exp_data and "experience" in exp_data) or
+                            ("condition" in exp_data and "experience" in exp_data)
+                    ):
+                        experiences.append(exp_data)
+
+                return experiences
+
+
+            # Handle single object
+            elif isinstance(parsed, dict) and (
+                    ("when_to_use" in parsed and "experience" in parsed) or
+                    ("condition" in parsed and "experience" in parsed)
+            ):
+                return [parsed]
+
+        # Fallback: try to parse entire response
+        parsed = json.loads(response)
+        if isinstance(parsed, list):
+            return parsed
+        elif isinstance(parsed, dict):
+            return [parsed]
+
+    except json.JSONDecodeError as e:
+        logger.warning(f"Failed to parse JSON experience response: {e}")
+
+    return []
+
+def get_trajectory_context(trajectory: Trajectory, step_sequence: List[Message]) -> str:
+    """Get context of step sequence within trajectory"""
+    try:
+        # Find position of step sequence in trajectory
+        start_idx = 0
+        for i, step in enumerate(trajectory.messages):
+            if step == step_sequence[0]:
+                start_idx = i
+                break
+
+        # Extract before and after context
+        context_before = trajectory.messages[max(0, start_idx - 2):start_idx]
+        context_after = trajectory.messages[start_idx + len(step_sequence):start_idx + len(step_sequence) + 2]
+
+        context = f"Query: {trajectory.metadata.get('query', 'N/A')}\n"
+
+        if context_before:
+            context += "Previous steps:\n" + "\n".join(
+                [f"- {step.content[:100]}..." for step in context_before]) + "\n"
+
+        if context_after:
+            context += "Following steps:\n" + "\n".join([f"- {step.content[:100]}..." for step in context_after])
+
+        return context
+
+    except Exception as e:
+        logger.error(f"Error getting trajectory context: {e}")
+        return f"Query: {trajectory.metadata.get('query', 'N/A')}"

llmflow/utils/registry.py

@@ -0,0 +1,33 @@
+from typing import List
+
+from loguru import logger
+
+from llmflow.utils.common_utils import camel_to_snake
+
+
+class Registry(object):
+    def __init__(self):
+        self._registry = {}
+
+    def register(self, name: str = ""):
+
+        def decorator(cls):
+            class_name = name if name else camel_to_snake(cls.__name__)
+            if class_name in self._registry:
+                logger.warning(f"name={class_name} is already registered, will be overwritten.")
+            self._registry[class_name] = cls
+            return cls
+
+        return decorator
+
+    def __getitem__(self, name: str):
+        if name not in self._registry:
+            raise KeyError(f"name={name} is not registered!")
+        return self._registry[name]
+
+    def __contains__(self, name: str):
+        return name in self._registry
+
+    @property
+    def registered_names(self) -> List[str]:
+        return sorted(self._registry.keys())

llmflow/utils/singleton.py

@@ -0,0 +1,9 @@
+def singleton(cls):
+    _instance = {}
+
+    def _singleton(*args, **kwargs):
+        if cls not in _instance:
+            _instance[cls] = cls(*args, **kwargs)
+        return _instance[cls]
+
+    return _singleton

llmflow/utils/timer.py

@@ -0,0 +1,53 @@
+import time
+
+from loguru import logger
+
+
+class Timer(object):
+    def __init__(self, name: str, use_ms: bool = False, stack_level: int = 2):
+        self.name: str = name
+        self.use_ms: bool = use_ms
+        self.stack_level: int = stack_level
+
+        self.time_start: float = 0
+        self.time_end: float = 0
+        self.time_cost: float = 0
+
+    def __enter__(self, *args, **kwargs):
+        self.time_start = time.time()
+        logger.info(f"---------- enter {self.name} ----------", stacklevel=self.stack_level)
+        return self
+
+    def __exit__(self, *args):
+        self.time_end = time.time()
+        self.time_cost = self.time_end - self.time_start
+        if self.use_ms:
+            time_str = f"{self.time_cost * 1000:.2f}ms"
+        else:
+            time_str = f"{self.time_cost:.3f}s"
+
+        logger.info(f"---------- leave {self.name} [{time_str}] ----------", stacklevel=self.stack_level)
+
+
+def timer(name: str = None, use_ms: bool = False, stack_level: int = 2):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            with Timer(name=name or func.__name__, use_ms=use_ms, stack_level=stack_level + 1):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+if __name__ == "__main__":
+    import random
+
+
+    @timer("run_func_final", use_ms=True)
+    def run_func():
+        time.sleep(random.uniform(0.05, 0.15))
+        print("done")
+
+
+    run_func()

llmflow/vector_store/__init__.py

@@ -0,0 +1,7 @@
+from llmflow.utils.registry import Registry
+
+VECTOR_STORE_REGISTRY = Registry()
+
+from llmflow.vector_store.es_vector_store import EsVectorStore
+from llmflow.vector_store.chroma_vector_store import ChromaVectorStore
+from llmflow.vector_store.file_vector_store import FileVectorStore