thinkai-framework 0.1.0__py3-none-any.whl

thinkai/__init__.py

@@ -0,0 +1,24 @@
+"""
+ThinkAi - Enterprise-grade AI Framework
+基于FastAPI的企业级AI大模型集成框架
+开箱即用,支持多模型,RAG,Agent等核心能力
+"""
+
+__version__ = "0.1.0"
+__author__ = "ThinkAi Team"
+
+from thinkai.core.client import ThinkAI
+from thinkai.core.config import Settings
+from thinkai.session.manager import SessionManager
+from thinkai.prompt.template import PromptTemplate
+from thinkai.rag.pipeline import RAGPipeline
+from thinkai.agent.base import Agent
+
+__all__ = [
+    "ThinkAI",
+    "Settings",
+    "SessionManager",
+    "PromptTemplate",
+    "RAGPipeline",
+    "Agent",
+]

thinkai/agent/__init__.py

@@ -0,0 +1,12 @@
+"""Agent模块"""
+from thinkai.agent.base import Agent, AgentConfig
+from thinkai.agent.tool import Tool, tool
+from thinkai.agent.react import ReActAgent
+
+__all__ = [
+    "Agent",
+    "AgentConfig",
+    "Tool",
+    "tool",
+    "ReActAgent",
+]

thinkai/agent/base.py

@@ -0,0 +1,83 @@
+"""Agent基类"""
+from abc import ABC, abstractmethod
+from typing import List, Optional, Dict, Any
+from pydantic import BaseModel, Field
+
+from thinkai.core.client import ThinkAI
+from thinkai.core.models import ChatMessage
+from thinkai.agent.tool import Tool
+
+
+class AgentConfig(BaseModel):
+    """Agent配置"""
+    name: str = Field(default="agent", description="Agent名称")
+    max_iterations: int = Field(default=10, description="最大迭代次数")
+    max_execution_time: int = Field(default=300, description="最大执行时间(秒)")
+    verbose: bool = Field(default=False, description="是否输出详细日志")
+
+
+class Agent(ABC):
+    """
+    Agent基类
+    """
+
+    def __init__(
+        self,
+        name: str = "agent",
+        model: Optional[str] = None,
+        tools: Optional[List[Tool]] = None,
+        max_iterations: int = 10,
+        verbose: bool = False,
+        ai_client: Optional[ThinkAI] = None,
+        **kwargs,
+    ):
+        self.name = name
+        self.model = model
+        self.tools = tools or []
+        self.max_iterations = max_iterations
+        self.verbose = verbose
+        self.ai_client = ai_client
+        self.extra = kwargs
+
+        # 注册工具
+        self._tool_registry: Dict[str, Tool] = {}
+        for tool in self.tools:
+            self._tool_registry[tool.name] = tool
+
+    @abstractmethod
+    async def run(self, task: str, **kwargs) -> str:
+        """
+        运行Agent
+        
+        Args:
+            task: 任务描述
+            **kwargs: 额外参数
+            
+        Returns:
+            执行结果
+        """
+        pass
+
+    def add_tool(self, tool: Tool):
+        """添加工具"""
+        self._tool_registry[tool.name] = tool
+        self.tools.append(tool)
+
+    def remove_tool(self, tool_name: str):
+        """移除工具"""
+        if tool_name in self._tool_registry:
+            del self._tool_registry[tool_name]
+            self.tools = [t for t in self.tools if t.name != tool_name]
+
+    def get_tool(self, tool_name: str) -> Optional[Tool]:
+        """获取工具"""
+        return self._tool_registry.get(tool_name)
+
+    def list_tools(self) -> List[Tool]:
+        """列出所有工具"""
+        return list(self._tool_registry.values())
+
+    def _log(self, message: str):
+        """日志输出"""
+        if self.verbose:
+            print(f"[{self.name}] {message}")

thinkai/agent/react.py

@@ -0,0 +1,148 @@
+"""ReAct Agent - 推理+行动循环"""
+from typing import Optional, List, Dict, Any
+import json
+import re
+from thinkai.core.client import ThinkAI
+from thinkai.core.models import ChatMessage
+from thinkai.agent.base import Agent
+from thinkai.agent.tool import Tool
+
+
+class ReActAgent(Agent):
+    """
+    ReAct Agent - 推理(Reasoning)和行动(Acting)交替执行
+    
+    工作流程:
+    1. 接收任务
+    2. 思考下一步行动
+    3. 调用工具(如果需要)
+    4. 观察结果
+    5. 重复2-4直到完成任务
+    """
+
+    def __init__(
+        self,
+        name: str = "react-agent",
+        model: Optional[str] = None,
+        tools: Optional[List[Tool]] = None,
+        max_iterations: int = 10,
+        verbose: bool = False,
+        ai_client: Optional[ThinkAI] = None,
+    ):
+        super().__init__(
+            name=name,
+            model=model,
+            tools=tools,
+            max_iterations=max_iterations,
+            verbose=verbose,
+            ai_client=ai_client,
+        )
+
+    async def run(self, task: str, **kwargs) -> str:
+        """
+        运行ReAct Agent
+        
+        Args:
+            task: 任务描述
+            **kwargs: 额外参数
+            
+        Returns:
+            最终结果
+        """
+        if not self.ai_client:
+            raise ValueError("ai_client is required for ReActAgent")
+        
+        # 构建工具描述
+        tools_desc = self._format_tools()
+        
+        # 构建系统提示
+        system_prompt = f"""You are a helpful AI assistant that can use tools to complete tasks.
+
+Available tools:
+{tools_desc}
+
+When you need to use a tool, respond with:
+Thought: [your reasoning]
+Action: [tool_name]
+Action Input: [tool input as JSON]
+
+After getting the tool result:
+Observation: [tool result]
+
+When you have enough information to answer:
+Thought: [your reasoning]
+Final Answer: [your final answer]
+
+Remember to think step by step and use tools when needed."""
+
+        messages = [
+            ChatMessage.system(system_prompt),
+            ChatMessage.user(task),
+        ]
+        
+        iteration = 0
+        
+        while iteration < self.max_iterations:
+            iteration += 1
+            self._log(f"Iteration {iteration}")
+            
+            # 获取LLM响应
+            response = await self.ai_client.chat(
+                messages=messages,
+                model=self.model,
+            )
+            
+            content = response.content or ""
+            messages.append(ChatMessage.assistant(content))
+            
+            self._log(f"Response: {content}")
+            
+            # 检查是否有Final Answer
+            if "Final Answer:" in content:
+                final_answer = content.split("Final Answer:")[-1].strip()
+                self._log(f"Final Answer: {final_answer}")
+                return final_answer
+            
+            # 解析Action
+            action_match = re.search(r"Action:\s*(\w+)", content)
+            action_input_match = re.search(r"Action Input:\s*(.+)", content, re.DOTALL)
+            
+            if action_match and action_input_match:
+                tool_name = action_match.group(1)
+                try:
+                    tool_input = json.loads(action_input_match.group(1).strip())
+                except json.JSONDecodeError:
+                    tool_input = {"input": action_input_match.group(1).strip()}
+                
+                self._log(f"Calling tool: {tool_name} with {tool_input}")
+                
+                # 执行工具
+                tool = self.get_tool(tool_name)
+                if not tool:
+                    observation = f"Error: Tool '{tool_name}' not found"
+                else:
+                    try:
+                        result = await tool.execute(**tool_input)
+                        observation = str(result)
+                    except Exception as e:
+                        observation = f"Error: {str(e)}"
+                
+                self._log(f"Observation: {observation}")
+                
+                messages.append(ChatMessage.user(f"Observation: {observation}"))
+            else:
+                # 没有Action,直接返回
+                return content
+        
+        return "Error: Max iterations reached"
+
+    def _format_tools(self) -> str:
+        """格式化工具描述"""
+        if not self.tools:
+            return "No tools available"
+        
+        tools_str = []
+        for tool in self.tools:
+            tools_str.append(f"- {tool.name}: {tool.description}")
+        
+        return "\n".join(tools_str)

thinkai/agent/tool.py

@@ -0,0 +1,115 @@
+"""工具定义"""
+from typing import Callable, Optional, Dict, Any, List
+import inspect
+import json
+from pydantic import BaseModel
+
+
+class Tool:
+    """
+    工具类 - 封装可被Agent调用的函数
+    
+    使用示例:
+        @tool
+        def search(query: str) -> str:
+            \"\"\"Search the web\"\"\"
+            return "results"
+        
+        或手动创建:
+        tool = Tool(
+            name="calculator",
+            description="Calculate math expressions",
+            func=calculate,
+            parameters={...}
+        )
+    """
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        func: Callable,
+        parameters: Optional[Dict[str, Any]] = None,
+    ):
+        self.name = name
+        self.description = description
+        self.func = func
+        self.parameters = parameters or self._infer_parameters(func)
+
+    def _infer_parameters(self, func: Callable) -> Dict[str, Any]:
+        """从函数签名推断参数schema"""
+        sig = inspect.signature(func)
+        properties = {}
+        required = []
+        
+        for name, param in sig.parameters.items():
+            param_type = "string"
+            if param.annotation == int:
+                param_type = "integer"
+            elif param.annotation == float:
+                param_type = "number"
+            elif param.annotation == bool:
+                param_type = "boolean"
+            elif param.annotation == list:
+                param_type = "array"
+            
+            properties[name] = {
+                "type": param_type,
+                "description": f"Parameter {name}",
+            }
+            
+            if param.default == inspect.Parameter.empty:
+                required.append(name)
+        
+        return {
+            "type": "object",
+            "properties": properties,
+            "required": required,
+        }
+
+    async def execute(self, **kwargs) -> Any:
+        """执行工具"""
+        if inspect.iscoroutinefunction(self.func):
+            return await self.func(**kwargs)
+        return self.func(**kwargs)
+
+    def to_openai_format(self) -> Dict[str, Any]:
+        """转换为OpenAI工具格式"""
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters,
+            },
+        }
+
+    def __repr__(self) -> str:
+        return f"Tool(name='{self.name}', description='{self.description}')"
+
+
+def tool(name: Optional[str] = None, description: Optional[str] = None):
+    """
+    工具装饰器
+    
+    使用示例:
+        @tool
+        def search(query: str) -> str:
+            \"\"\"Search the web\"\"\"
+            ...
+        
+        @tool(name="calc", description="Calculate")
+        def calculate(expr: str) -> str:
+            ...
+    """
+    def decorator(func: Callable) -> Tool:
+        tool_name = name or func.__name__
+        tool_desc = description or (func.__doc__ or "").strip()
+        
+        return Tool(
+            name=tool_name,
+            description=tool_desc,
+            func=func,
+        )
+    
+    return decorator

thinkai/core/__init__.py

@@ -0,0 +1,22 @@
+"""核心模块"""
+from thinkai.core.config import Settings
+from thinkai.core.client import ThinkAI
+from thinkai.core.models import (
+    ChatMessage,
+    ChatRequest,
+    ChatResponse,
+    ChatChoice,
+    Usage,
+    MessageRole,
+)
+
+__all__ = [
+    "Settings",
+    "ThinkAI",
+    "ChatMessage",
+    "ChatRequest",
+    "ChatResponse",
+    "ChatChoice",
+    "Usage",
+    "MessageRole",
+]