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

llmflow/tool/tavily_search_tool.py

@@ -1,109 +0,0 @@
-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

@@ -1,23 +0,0 @@
-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

@@ -1,17 +0,0 @@
-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

@@ -1,25 +0,0 @@
-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

@@ -1,156 +0,0 @@
-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

@@ -1,102 +0,0 @@
-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

@@ -1,33 +0,0 @@
-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/vector_store/__init__.py

@@ -1,7 +0,0 @@
-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

llmflow/vector_store/base_vector_store.py

@@ -1,136 +0,0 @@
-import fcntl
-import json
-from abc import ABC
-from pathlib import Path
-from typing import List, Iterable
-
-from loguru import logger
-from pydantic import BaseModel, Field
-from tqdm import tqdm
-
-from llmflow.embedding_model.base_embedding_model import BaseEmbeddingModel
-from llmflow.schema.vector_node import VectorNode
-
-
-class BaseVectorStore(BaseModel, ABC):
-    embedding_model: BaseEmbeddingModel | None = Field(default=None)
-    batch_size: int = Field(default=1024)
-
-    @staticmethod
-    def _load_from_path(workspace_id: str, path: str | Path, callback_fn=None, **kwargs) -> Iterable[VectorNode]:
-        workspace_path = Path(path) / f"{workspace_id}.jsonl"
-        if not workspace_path.exists():
-            logger.warning(f"workspace_path={workspace_path} is not exists!")
-            return
-
-        with workspace_path.open() as f:
-            fcntl.flock(f, fcntl.LOCK_SH)
-            try:
-                for line in tqdm(f, desc="load from path"):
-                    if line.strip():
-                        node_dict = json.loads(line.strip())
-                        if callback_fn:
-                            node = callback_fn(node_dict)
-                        else:
-                            node = VectorNode(**node_dict, **kwargs)
-                        node.workspace_id = workspace_id
-                        yield node
-
-            finally:
-                fcntl.flock(f, fcntl.LOCK_UN)
-
-    @staticmethod
-    def _dump_to_path(nodes: Iterable[VectorNode], workspace_id: str, path: str | Path = "", callback_fn=None,
-                      ensure_ascii: bool = False, **kwargs):
-        dump_path: Path = Path(path)
-        dump_path.mkdir(parents=True, exist_ok=True)
-        dump_file = dump_path / f"{workspace_id}.jsonl"
-
-        count = 0
-        with dump_file.open("w") as f:
-            fcntl.flock(f, fcntl.LOCK_EX)
-            try:
-                for node in tqdm(nodes, desc="dump to path"):
-                    node.workspace_id = workspace_id
-                    if callback_fn:
-                        node_dict = callback_fn(node)
-                    else:
-                        node_dict = node.model_dump()
-                    assert isinstance(node_dict, dict)
-                    f.write(json.dumps(node_dict, ensure_ascii=ensure_ascii, **kwargs))
-                    f.write("\n")
-                    count += 1
-
-                return {"size": count}
-            finally:
-                fcntl.flock(f, fcntl.LOCK_UN)
-
-    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):
-        if not self.exist_workspace(workspace_id=workspace_id, **kwargs):
-            logger.warning(f"workspace_id={workspace_id} is not exist!")
-            return {}
-
-        return self._dump_to_path(nodes=self._iter_workspace_nodes(workspace_id=workspace_id, **kwargs),
-                                  workspace_id=workspace_id,
-                                  path=path,
-                                  callback_fn=callback_fn,
-                                  **kwargs)
-
-    def load_workspace(self, workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, callback_fn=None,
-                       **kwargs):
-        if self.exist_workspace(workspace_id, **kwargs):
-            self.delete_workspace(workspace_id=workspace_id, **kwargs)
-            logger.info(f"delete workspace_id={workspace_id}")
-
-        self.create_workspace(workspace_id=workspace_id, **kwargs)
-
-        all_nodes: List[VectorNode] = []
-        if nodes:
-            all_nodes.extend(nodes)
-        for node in self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs):
-            all_nodes.append(node)
-        self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs)
-        return {"size": len(all_nodes)}
-
-    def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs):
-        if not self.exist_workspace(workspace_id=src_workspace_id, **kwargs):
-            logger.warning(f"src_workspace_id={src_workspace_id} is not exist!")
-            return {}
-
-        if not self.exist_workspace(dest_workspace_id, **kwargs):
-            self.create_workspace(workspace_id=dest_workspace_id, **kwargs)
-
-        nodes = []
-        node_size = 0
-        for node in self._iter_workspace_nodes(workspace_id=src_workspace_id, **kwargs):
-            nodes.append(node)
-            node_size += 1
-            if len(nodes) >= self.batch_size:
-                self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs)
-                nodes.clear()
-
-        if nodes:
-            self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs)
-        return {"size": node_size}
-
-    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
-