flowllm 0.1.0__py3-none-any.whl

llmflow/schema/experience.py

@@ -0,0 +1,144 @@
+import datetime
+from abc import ABC
+from typing import List
+from uuid import uuid4
+
+from loguru import logger
+from pydantic import BaseModel, Field
+
+from llmflow.schema.vector_node import VectorNode
+
+
+class ExperienceMeta(BaseModel):
+    author: str = Field(default="")
+    created_time: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
+    modified_time: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
+    extra_info: dict | None = Field(default=None)
+
+    def update_modified_time(self):
+        self.modified_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+
+class BaseExperience(BaseModel, ABC):
+    workspace_id: str = Field(default="")
+
+    experience_id: str = Field(default_factory=lambda: uuid4().hex)
+    experience_type: str = Field(default="")
+
+    when_to_use: str = Field(default="")
+    content: str | bytes = Field(default="")
+    score: float | None = Field(default=None)
+    metadata: ExperienceMeta = Field(default_factory=ExperienceMeta)
+
+    def to_vector_node(self) -> VectorNode:
+        raise NotImplementedError
+
+    @classmethod
+    def from_vector_node(cls, node: VectorNode):
+        raise NotImplementedError
+
+
+class TextExperience(BaseExperience):
+    experience_type: str = Field(default="text")
+
+    def to_vector_node(self) -> VectorNode:
+        return VectorNode(unique_id=self.experience_id,
+                          workspace_id=self.workspace_id,
+                          content=self.when_to_use,
+                          metadata={
+                              "experience_type": self.experience_type,
+                              "experience_content": self.content,
+                              "score": self.score,
+                              "metadata": self.metadata.model_dump(),
+                          })
+
+    @classmethod
+    def from_vector_node(cls, node: VectorNode):
+        return cls(workspace_id=node.workspace_id,
+                   experience_id=node.unique_id,
+                   experience_type=node.metadata.get("experience_type"),
+                   when_to_use=node.content,
+                   content=node.metadata.get("experience_content"),
+                   score=node.metadata.get("score"),
+                   metadata=node.metadata.get("metadata"))
+
+
+class FunctionArg(BaseModel):
+    arg_name: str = Field(default=...)
+    arg_type: str = Field(default=...)
+    required: bool = Field(default=True)
+
+
+class Function(BaseModel):
+    func_code: str = Field(default=..., description="function code")
+    func_name: str = Field(default=..., description="function name")
+    func_args: List[FunctionArg] = Field(default_factory=list)
+
+
+class FuncExperience(BaseExperience):
+    experience_type: str = Field(default="function")
+    functions: List[Function] = Field(default_factory=list)
+
+
+class PersonalExperience(BaseExperience):
+    experience_type: str = Field(default="personal")
+    person: str = Field(default="")
+    topic: str = Field(default="")
+
+
+class KnowledgeExperience(BaseExperience):
+    experience_type: str = Field(default="knowledge")
+    topic: str = Field(default="")
+
+
+def vector_node_to_experience(node: VectorNode) -> BaseExperience:
+    experience_type = node.metadata.get("experience_type")
+    if experience_type == "text":
+        return TextExperience.from_vector_node(node)
+
+    elif experience_type == "function":
+        return FuncExperience.from_vector_node(node)
+
+    elif experience_type == "personal":
+        return PersonalExperience.from_vector_node(node)
+
+    elif experience_type == "knowledge":
+        return KnowledgeExperience.from_vector_node(node)
+
+    else:
+        logger.warning(f"experience type {experience_type} not supported")
+        return TextExperience.from_vector_node(node)
+
+
+def dict_to_experience(experience_dict: dict) -> BaseExperience:
+    experience_type = experience_dict.get("experience_type", "text")
+    if experience_type == "text":
+        return TextExperience(**experience_dict)
+
+    elif experience_type == "function":
+        return FuncExperience(**experience_dict)
+
+    elif experience_type == "personal":
+        return PersonalExperience(**experience_dict)
+
+    elif experience_type == "knowledge":
+        return KnowledgeExperience(**experience_dict)
+
+    else:
+        logger.warning(f"experience type {experience_type} not supported")
+        return TextExperience(**experience_dict)
+
+
+if __name__ == "__main__":
+    e1 = TextExperience(
+        workspace_id="w_1024",
+        experience_id="123",
+        when_to_use="test case use",
+        content="test content",
+        score=0.99,
+        metadata=ExperienceMeta(author="user"))
+    print(e1.model_dump_json(indent=2))
+    v1 = e1.to_vector_node()
+    print(v1.model_dump_json(indent=2))
+    e2 = vector_node_to_experience(v1)
+    print(e2.model_dump_json(indent=2))

llmflow/schema/message.py

@@ -0,0 +1,68 @@
+import json
+from typing import List
+
+from pydantic import BaseModel, Field, model_validator
+
+from llmflow.enumeration.role import Role
+
+
+class ToolCall(BaseModel):
+    index: int = Field(default=0)
+    id: str = Field(default="")
+    name: str = Field(default="")
+    arguments: str = Field(default="")
+    type: str = Field(default="function")
+
+    @model_validator(mode="before")  # noqa
+    @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)
+
+    def simple_dump(self) -> dict:
+        return {
+            "id": self.id,
+            self.type: {
+                "arguments": self.arguments,
+                "name": self.name
+            },
+            "type": self.type,
+            "index": self.index,
+        }
+
+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_when_empty: bool = True) -> dict:
+        result: dict
+        if self.content:
+            result = {"role": self.role.value, "content": self.content}
+        elif add_reason_when_empty 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_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)

llmflow/schema/request.py

@@ -0,0 +1,32 @@
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from llmflow.schema.message import Message, Trajectory
+
+
+class BaseRequest(BaseModel):
+    workspace_id: str = Field(default="default")
+    config: dict = Field(default_factory=dict)
+
+
+class RetrieverRequest(BaseRequest):
+    query: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)
+    top_k: int = Field(default=1)
+
+
+class SummarizerRequest(BaseRequest):
+    traj_list: List[Trajectory] = Field(default_factory=list)
+
+
+class VectorStoreRequest(BaseRequest):
+    action: str = Field(default="")
+    src_workspace_id: str = Field(default="")
+    path: str = Field(default="")
+
+
+class AgentRequest(BaseRequest):
+    query: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)
+

llmflow/schema/response.py

@@ -0,0 +1,29 @@
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from llmflow.schema.experience import BaseExperience
+from llmflow.schema.message import Message
+
+
+class BaseResponse(BaseModel):
+    success: bool = Field(default=True)
+    metadata: dict = Field(default_factory=dict)
+
+
+class RetrieverResponse(BaseResponse):
+    experience_list: List[BaseExperience] = Field(default_factory=list)
+    experience_merged: str = Field(default="")
+
+
+class SummarizerResponse(BaseResponse):
+    experience_list: List[BaseExperience] = Field(default_factory=list)
+    deleted_experience_ids: List[str] = Field(default_factory=list)
+
+
+class VectorStoreResponse(BaseResponse):
+    ...
+
+class AgentResponse(BaseResponse):
+    answer: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)

llmflow/schema/vector_node.py

@@ -0,0 +1,11 @@
+from typing import List
+from uuid import uuid4
+from pydantic import BaseModel, Field
+
+
+class VectorNode(BaseModel):
+    unique_id: str = Field(default_factory=lambda: uuid4().hex)
+    workspace_id: str = Field(default="")
+    content: str = Field(default="")
+    vector: List[float] | None = Field(default=None)
+    metadata: dict = Field(default_factory=dict)

llmflow/service/llmflow_service.py

@@ -0,0 +1,96 @@
+from concurrent.futures import ThreadPoolExecutor
+from typing import List
+
+from loguru import logger
+
+from llmflow.config.config_parser import ConfigParser
+from llmflow.embedding_model import EMBEDDING_MODEL_REGISTRY
+from llmflow.pipeline.pipeline import Pipeline
+from llmflow.pipeline.pipeline_context import PipelineContext
+from llmflow.schema.app_config import AppConfig, HttpServiceConfig, EmbeddingModelConfig
+from llmflow.schema.request import SummarizerRequest, RetrieverRequest, VectorStoreRequest, AgentRequest, \
+    BaseRequest
+from llmflow.schema.response import SummarizerResponse, RetrieverResponse, VectorStoreResponse, AgentResponse, \
+    BaseResponse
+from llmflow.vector_store import VECTOR_STORE_REGISTRY
+
+
+class LLMFlowService:
+
+    def __init__(self, args: List[str]):
+        self.config_parser = ConfigParser(args)
+        self.init_app_config: AppConfig = self.config_parser.get_app_config()
+        self.thread_pool = ThreadPoolExecutor(max_workers=self.init_app_config.thread_pool.max_workers)
+
+        # The vectorstore is initialized at the very beginning and then used directly afterward.
+        self.vector_store_dict: dict = {}
+        for name, config in self.init_app_config.vector_store.items():
+            assert config.backend in VECTOR_STORE_REGISTRY, f"backend={config.backend} is not existed"
+            vector_store_cls = VECTOR_STORE_REGISTRY[config.backend]
+
+            assert config.embedding_model in self.init_app_config.embedding_model, \
+                f"embedding_model={config.embedding_model} is not existed"
+            embedding_model_config: EmbeddingModelConfig = self.init_app_config.embedding_model[config.embedding_model]
+
+            assert embedding_model_config.backend in EMBEDDING_MODEL_REGISTRY, \
+                f"embedding_model={embedding_model_config.backend} is not existed"
+            embedding_model_cls = EMBEDDING_MODEL_REGISTRY[embedding_model_config.backend]
+            embedding_model = embedding_model_cls(model_name=embedding_model_config.model_name,
+                                                  **embedding_model_config.params)
+
+            self.vector_store_dict[name] = vector_store_cls(embedding_model=embedding_model, **config.params)
+
+    @property
+    def http_service_config(self) -> HttpServiceConfig:
+        return self.init_app_config.http_service
+
+    def __call__(self, api: str, request: dict | BaseRequest) -> BaseResponse:
+        if isinstance(request, dict):
+            app_config: AppConfig = self.config_parser.get_app_config(**request["config"])
+        else:
+            app_config: AppConfig = self.config_parser.get_app_config(**request.config)
+
+        if api == "retriever":
+            if isinstance(request, dict):
+                request = RetrieverRequest(**request)
+            response = RetrieverResponse()
+            pipeline = app_config.api.retriever
+
+        elif api == "summarizer":
+            if isinstance(request, dict):
+                request = SummarizerRequest(**request)
+            response = SummarizerResponse()
+            pipeline = app_config.api.summarizer
+
+        elif api == "vector_store":
+            if isinstance(request, dict):
+                request = VectorStoreRequest(**request)
+            response = VectorStoreResponse()
+            pipeline = app_config.api.vector_store
+
+        elif api == "agent":
+            if isinstance(request, dict):
+                request = AgentRequest(**request)
+            response = AgentResponse()
+            pipeline = app_config.api.agent
+
+        else:
+            raise RuntimeError(f"Invalid service.api={api}")
+
+        logger.info(f"request={request.model_dump_json()}")
+
+        try:
+            context = PipelineContext(app_config=app_config,
+                                      thread_pool=self.thread_pool,
+                                      request=request,
+                                      response=response,
+                                      vector_store_dict=self.vector_store_dict)
+            pipeline = Pipeline(pipeline=pipeline, context=context)
+            pipeline()
+
+        except Exception as e:
+            logger.exception(f"api={api} encounter error={e.args}")
+            response.success = False
+            response.metadata["error"] = str(e)
+
+        return response

llmflow/tool/__init__.py

@@ -0,0 +1,9 @@
+from llmflow.utils.registry import Registry
+
+TOOL_REGISTRY = Registry()
+
+from llmflow.tool.code_tool import CodeTool
+from llmflow.tool.dashscope_search_tool import DashscopeSearchTool
+from llmflow.tool.tavily_search_tool import TavilySearchTool
+from llmflow.tool.terminate_tool import TerminateTool
+from llmflow.tool.mcp_tool import MCPTool

llmflow/tool/base_tool.py

@@ -0,0 +1,80 @@
+from abc import ABC
+
+from loguru import logger
+from pydantic import BaseModel, Field
+
+
+class BaseTool(BaseModel, ABC):
+    tool_id: str = Field(default="")
+    name: str = Field(..., description="tool name")
+    description: str = Field(..., description="tool description")
+    tool_type: str = Field(default="function")
+    parameters: dict = Field(default_factory=dict, description="tool parameters")
+    arguments: dict = Field(default_factory=dict, description="execute arguments")
+
+    enable_cache: bool = Field(default=False, description="whether to cache the tool result")
+    cached_result: dict = Field(default_factory=dict, description="tool execution result")
+
+    max_retries: int = Field(default=3, description="max retries")
+    raise_exception: bool = Field(default=True, description="raise exception")
+    success: bool = Field(default=True, description="whether the tool executed successfully")
+
+    def reset(self):
+        self.arguments.clear()
+        self.cached_result.clear()
+        self.success = True
+
+    def _execute(self, **kwargs):
+        raise NotImplementedError
+
+    def execute(self, **kwargs):
+        cache_id = ""
+        if self.enable_cache:
+            cache_id = self.get_cache_id(**kwargs)
+            if cache_id in self.cached_result:
+                return self.cached_result[cache_id]
+
+        for i in range(self.max_retries):
+            try:
+                if self.enable_cache:
+                    self.cached_result[cache_id] = self._execute(**kwargs)
+                    return self.cached_result[cache_id]
+
+                else:
+                    return self._execute(**kwargs)
+
+            except Exception as e:
+                logger.exception(f"using tool.name={self.name} encounter error with e={e.args}")
+                if i == self.max_retries - 1 and self.raise_exception:
+                    raise e
+
+        return None
+
+
+    def simple_dump(self) -> dict:
+        """
+        It may be in other different tool params formats; different versions are completed here.
+        """
+        return {
+            "type": self.tool_type,
+            self.tool_type: {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters,
+            },
+        }
+
+    @property
+    def input_schema(self) -> dict:
+        return self.parameters.get("properties", {})
+
+    @property
+    def output_schema(self) -> dict:
+        raise NotImplementedError
+
+    def refresh(self):
+        # for mcp
+        raise NotImplementedError
+
+    def get_cache_id(self, **kwargs) -> str:
+        raise NotImplementedError

llmflow/tool/code_tool.py

@@ -0,0 +1,43 @@
+import sys
+from io import StringIO
+
+from llmflow.tool import TOOL_REGISTRY
+from llmflow.tool.base_tool import BaseTool
+
+
+@TOOL_REGISTRY.register()
+class CodeTool(BaseTool):
+    name: str = "python_execute"
+    description: str = "Execute python code can be used in scenarios such as analysis or calculation, and the final result can be printed using the `print` function."
+    parameters: dict = {
+        "type": "object",
+        "properties": {
+            "code": {
+                "type": "string",
+                "description": "code to be executed. Please do not execute any matplotlib code here.",
+            }
+        },
+        "required": ["code"]
+    }
+
+    def _execute(self, code: str, **kwargs):
+        old_stdout = sys.stdout
+        redirected_output = sys.stdout = StringIO()
+
+        try:
+            exec(code)
+            result = redirected_output.getvalue()
+
+        except Exception as e:
+            self.success = False
+            result = str(e)
+
+        sys.stdout = old_stdout
+
+        return result
+
+
+if __name__ == '__main__':
+    tool = CodeTool()
+    print(tool.execute(code="print('Hello World')"))
+    print(tool.execute(code="print('Hello World!'"))

llmflow/tool/dashscope_search_tool.py

@@ -0,0 +1,162 @@
+import os
+from typing import Literal
+
+import dashscope
+from dashscope.api_entities.dashscope_response import Message
+from dotenv import load_dotenv
+from loguru import logger
+from pydantic import Field
+
+from llmflow.tool import TOOL_REGISTRY
+from llmflow.tool.base_tool import BaseTool
+
+
+@TOOL_REGISTRY.register()
+class DashscopeSearchTool(BaseTool):
+    name: str = "web_search"
+    description: str = "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."
+    parameters: dict = {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "search keyword",
+            }
+        },
+        "required": ["query"]
+    }
+
+    model_name: Literal["qwen-plus-2025-04-28", "qwq-plus-latest", "qwen-max-2025-01-25"] = \
+        Field(default="qwen-plus-2025-04-28")
+    api_key: str = Field(default_factory=lambda: os.environ["DASHSCOPE_API_KEY"])
+    stream_print: bool = Field(default=False)
+    temperature: float = Field(default=0.0000001)
+    use_role_prompt: bool = Field(default=True)
+    role_prompt: str = """
+# user's question
+{question}
+
+# task
+Extract the original content related to the user's question directly from the context, maintain accuracy, and avoid excessive processing.    """.strip()
+    return_only_content: bool = Field(default=True)
+
+    def parse_reasoning_response(self, response, result: dict):
+        is_answering = False
+        is_first_chunk = True
+
+        for chunk in response:
+            if is_first_chunk:
+                result["search_results"] = chunk.output.search_info["search_results"]
+
+                if self.stream_print:
+                    print("=" * 20 + "search result" + "=" * 20)
+                    for web in result["search_results"]:
+                        print(f"[{web['index']}]: [{web['title']}]({web['url']})")
+                    print("=" * 20 + "thinking process" + "=" * 20)
+                result["reasoning_content"] += chunk.output.choices[0].message.reasoning_content
+
+                if self.stream_print:
+                    print(chunk.output.choices[0].message.reasoning_content, end="", flush=True)
+                is_first_chunk = False
+
+            else:
+                if chunk.output.choices[0].message.content == "" \
+                        and chunk.output.choices[0].message.reasoning_content == "":
+                    pass
+
+                else:
+                    if chunk.output.choices[0].message.reasoning_content != "" and \
+                            chunk.output.choices[0].message.content == "":
+
+                        if self.stream_print:
+                            print(chunk.output.choices[0].message.reasoning_content, end="", flush=True)
+                        result["reasoning_content"] += chunk.output.choices[0].message.reasoning_content
+
+                    elif chunk.output.choices[0].message.content != "":
+                        if not is_answering:
+                            if self.stream_print:
+                                print("\n" + "=" * 20 + "complete answer" + "=" * 20)
+                            is_answering = True
+
+                        if self.stream_print:
+                            print(chunk.output.choices[0].message.content, end="", flush=True)
+                        result["answer_content"] += chunk.output.choices[0].message.content
+
+    def parse_response(self, response, result: dict):
+        is_first_chunk = True
+
+        for chunk in response:
+            if is_first_chunk:
+                result["search_results"] = chunk.output.search_info["search_results"]
+
+                if self.stream_print:
+                    print("=" * 20 + "search result" + "=" * 20)
+                    for web in result["search_results"]:
+                        print(f"[{web['index']}]: [{web['title']}]({web['url']})")
+                    print("\n" + "=" * 20 + "complete answer" + "=" * 20)
+                is_first_chunk = False
+
+            else:
+                if chunk.output.choices[0].message.content == "":
+                    pass
+
+                else:
+                    if chunk.output.choices[0].message.content != "":
+                        if self.stream_print:
+                            print(chunk.output.choices[0].message.content, end="", flush=True)
+                        result["answer_content"] += chunk.output.choices[0].message.content
+
+    def execute(self, query: str = "", **kwargs):
+        result = {
+            "search_results": [],
+            "reasoning_content": "",
+            "answer_content": ""
+        }
+        user_query = self.role_prompt.format(question=query) if self.use_role_prompt else query
+        messages = [Message(role="user", content=user_query)]
+
+        response = dashscope.Generation.call(
+            api_key=self.api_key,
+            model=self.model_name,
+            messages=messages,
+            enable_thinking=True,
+            enable_search=True,
+            search_options={
+                "forced_search": True,
+                "enable_source": True,
+                "enable_citation": False,
+                "search_strategy": "pro"
+            },
+            stream=True,
+            incremental_output=True,
+            result_format="message",
+        )
+
+        if self.model_name != "qwen-max-2025-01-25":
+            self.parse_reasoning_response(response, result)
+        else:
+            self.parse_response(response, result)
+
+        if self.return_only_content:
+            return result["answer_content"]
+        else:
+            return result
+
+
+def main():
+    load_dotenv()
+    query = "What is artificial intelligence?"
+
+    tool = DashscopeSearchTool(stream_print=True)
+    logger.info(tool.execute(query=query))
+
+    tool = DashscopeSearchTool(stream_print=False)
+    logger.info(tool.execute(query=query))
+
+    tool = DashscopeSearchTool(stream_print=True, model_name="qwen-max-2025-01-25")
+    logger.info(tool.execute(query=query))
+
+
+if __name__ == '__main__':
+    main()