agent = FunctionAgent(
    tools=[multiply, add, subtract, divide],
    llm=Ollama(model="llama3.1", request_timeout=360.0),
    system_prompt="You are a math assistant and must prioritize using the provided tool functions (add/subtract/multiply/divide) to complete calculations. Direct calculation using your own abilities is prohibited. For each calculation request, you must call the corresponding tool function and then return a clear calculation result.",
)

import asyncio
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.ollama import Ollama

# 让AI调用multiply方法，但是未使用历史功能
def multiply(a: float, b: float) -> float:
    """Multiply two numbers and return the product."""
    print(f"multiply方法执行({a}, {b})")
    return a * b

def add(a: float, b: float) -> float:
    """Add two numbers and return the sum."""
    print(f"add方法执行({a}, {b})")
    return a + b

def subtract(a: float, b: float) -> float:
    """Subtract the second number from the first and return the difference."""
    print(f"subtract方法执行({a}, {b})")
    return a - b

def divide(a: float, b: float) -> float:
    """Divide the first number by the second and return the quotient.
    Raises ValueError if the divisor (b) is zero.
    """
    print(f"divide方法执行({a}, {b})")
    if b == 0:
        raise ValueError("除法运算中，除数不能为0！")
    return a / b

agent = FunctionAgent(
    tools=[multiply, add, subtract, divide],
    llm=Ollama(model="llama3.1", request_timeout=360.0),
    system_prompt="You are a math assistant and must prioritize using the provided tool functions (add/subtract/multiply/divide) to complete calculations. Direct calculation using your own abilities is prohibited. For each calculation request, you must call the corresponding tool function and then return a clear calculation result.",
)

async def main():
    response = await agent.run("7+8 and 9*3")
    print(str(response))

if __name__ == "__main__":
    asyncio.run(main())

query = "作者大学学的什么专业？另外，7*8等于多少？"

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
import asyncio
import os

if os.name == 'nt':
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

Settings.embed_model = OllamaEmbedding(
    model_name="nomic-embed-text:latest",
    base_url="http://localhost:11434",
)

Settings.llm = Ollama(
    model="llama3.1",
    base_url="http://localhost:11434",
    request_timeout=360.0,
    verbose=True  
)

data_dir = "data"
if not os.path.exists(data_dir):
    print(f"{data_dir}directory no found")
    exit(1)

documents = SimpleDirectoryReader(data_dir).load_data()
print(documents)
print(f"加载 {len(documents)} 个文档")
if len(documents) == 0:
    print("No available information detected in the data folder")
    exit(1)

index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine(
    top_k=3,  
    verbose=True  
)

def multiply(a: float, b: float) -> float:
    """Multiply two numbers and return the product."""
    print(f"multiply方法执行({a}, {b})")
    return a * b

def add(a: float, b: float) -> float:
    """Add two numbers and return the sum."""
    print(f"add方法执行({a}, {b})")
    return a + b

def subtract(a: float, b: float) -> float:
    """Subtract the second number from the first and return the difference."""
    print(f"subtract方法执行({a}, {b})")
    return a - b

def divide(a: float, b: float) -> float:
    """Divide the first number by the second and return the quotient.
    Raises ValueError if the divisor (b) is zero.
    """
    print(f"divide方法执行({a}, {b})")
    if b == 0:
        raise ValueError("除法运算中，除数不能为0！")
    return a / b

async def search_documents(query: str) -> str:
    """Search documents and return the results."""
    print(f"search_documents function called-> {query}") 
    response = await query_engine.aquery(query)
    return f"文档检索结果：{str(response)}"

agent = AgentWorkflow.from_tools_or_functions(
    [multiply, search_documents, add, subtract, divide],
    llm=Settings.llm,
    system_prompt="""
    You are a helpful assistant that can perform calculations
    and search through documents to answer questions.
    """, 
    verbose=True  
)

async def main():
    query = "作者大学学的什么专业？另外，7*8等于多少？"
    print(f"提问：{query}")
    response = await agent.run(query)
    print("\n===== 结果 =====")
    print(response)

if __name__ == "__main__":
    asyncio.run(main())

query = "作者的高中母校是哪里？另外，计算8除以0"

def divide(a: float, b: float) -> float:
    """[Mathematical Tool] Division operation, only accepts integer/float parameters, divisor cannot be 0"""
    if not isinstance(a, (int, float)) or not isinstance(b, (int, float)):
        raise TypeError(f"Division tool only accepts numeric parameters! Invalid current parameters: a={a}, b={b}")
    print(f"divide方法执行({a}, {b})")
    if b == 0:
        raise ValueError("【Division Exception】Divisor cannot be 0, cannot calculate 8÷0!")
    return a / b

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.core.workflow import Context
from llama_index.core.tools import FunctionTool  # 显式封装工具，精准控用途
import asyncio
import os

if os.name == 'nt':
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

# 模型配置
Settings.embed_model = OllamaEmbedding(
    model_name="nomic-embed-text:latest",
    base_url="http://localhost:11434",
)

Settings.llm = Ollama(
    model="llama3.1",
    base_url="http://localhost:11434",
    request_timeout=360.0,
    verbose=True  
)

# 文档加载
data_dir = "data"
if not os.path.exists(data_dir):
    print(f"{data_dir} directory not found")
    exit(1)

documents = SimpleDirectoryReader(data_dir).load_data()
print(f"加载 {len(documents)} 个文档")
if len(documents) == 0:
    print("No available information detected in the data folder")
    exit(1)

index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine(top_k=3, verbose=True)

def add(a: float, b: float) -> float:
    """[Mathematical Tool] Addition operation, only accepts integer/float parameters"""
    if not isinstance(a, (int, float)) or not isinstance(b, (int, float)):
        raise TypeError(f"The addition tool only accepts numeric parameters! The current parameter is invalid: a={a}, b={b}")
    print(f"add方法执行({a}, {b})")
    return a + b

def subtract(a: float, b: float) -> float:
    """[Mathematical Tool] Subtraction operation, calculate the difference of a-b, where parameter a is the minuend and b is the subtrahend."""
    print(f"subtract方法执行({a}, {b})")
    return a - b

def multiply(a: float, b: float) -> float:
    """[Mathematical Tool] Multiplication operation, calculate the product of a*b, where parameter a is the multiplicand and b is the multiplier."""
    print(f"multiply方法执行({a}, {b})")
    return a * b

def divide(a: float, b: float) -> float:
    """[Mathematical Tool] Division operation, only accepts integer/float parameters, divisor cannot be 0"""
    if not isinstance(a, (int, float)) or not isinstance(b, (int, float)):
        raise TypeError(f"Division tool only accepts numeric parameters! Invalid current parameters: a={a}, b={b}")
    print(f"divide方法执行({a}, {b})")
    if b == 0:
        raise ValueError("【Division Exception】Divisor cannot be 0, cannot calculate 8÷0!")
    return a / b

async def search_documents(query: str) -> str:
    """[Document Tool] Used only for retrieving information from local documents, such as 'the author's university major.'  
    This tool must not be used for non-document questions (such as names or calculations)!  
    If there are no results, return the standardized prompt.
    """
    print(f"search_documents function called-> {query}") 
    response = await query_engine.aquery(query)
    if not str(response).strip() or "未找到" in str(response):
        return "【Document Exception】No relevant information found in the documents."
    return f"文档检索结果：{str(response)}"

# ========== 2. 显式封装工具：给LLM明确的工具用途标签 ==========
math_tools = [
    FunctionTool.from_defaults(fn=add),
    FunctionTool.from_defaults(fn=subtract),
    FunctionTool.from_defaults(fn=multiply),
    FunctionTool.from_defaults(fn=divide)
]
doc_tool = [FunctionTool.from_defaults(fn=search_documents)]

# ========== 3. 核心：Prompt 精准划分任务边界（重点！） ==========
agent = AgentWorkflow.from_tools_or_functions(
    math_tools + doc_tool,  # 数学工具放前面，优先调用
    llm=Settings.llm,
    system_prompt="""
You are a [Document Retrieval & Math Calculation] assistant with memory capabilities. Follow the clear steps below and do not make comments about rules:

1. Memory Tasks (use your brain directly, do not use any tools)
- Applicable scenarios: storing names, asking for names, instructions like "remember my name";
- Operation: read the Context (this is your memory carrier, not a tool), respond directly, do not mention "violating rules";
- Prohibited: calling any tools like search_documents/add.

2. Document Retrieval Tasks (must be precise, no shortcuts)
- Applicable scenarios: questions like "author's high school alma mater" or anything related to the author;
- Operation: the search term must **exactly match** the user's question (for example, if the user asks "Where is the author's high school alma mater," use this as the search term);
- Prohibited: simplifying the search term (e.g., just using author), or guessing if the document has no information.

3. Math Calculation Tasks (must use tools, do not calculate manually)
- Applicable scenarios: addition, subtraction, multiplication, division (e.g., 8÷0);
- Operation: strictly use the corresponding tools (for division, use divide), if the tool throws an exception, return the exception message directly;
- Prohibited: do not calculate manually and respond with "undefined."

4. Multi-Intent Questions: split according to the above three types of tasks, handle separately, and then integrate the results.
""",
    verbose=True  
)

# ========== 4. 多轮对话测试（全中文提问，减少解析混乱） ==========
async def main():
    ctx = Context(agent)  # 记忆容器

    # 第一轮：告知名字（记忆类，不调用工具）
    print("===== 第一轮对话 =====")
    query1 = "My name is Lihua"
    print(f"提问：{query1}")
    response1 = await agent.run(query1, ctx=ctx)
    print(f"回复：{response1}\n")

    # 第二轮：问名字（记忆类，不调用工具）
    print("===== 第二轮对话 =====")
    query2 = "What is my name?"
    print(f"提问：{query2}")
    response2 = await agent.run(query2, ctx=ctx)
    print(f"回复：{response2}\n")

    # 第三轮：文档+计算（调用对应工具，同时记忆名字）
    print("===== 第三轮对话 =====")
    query3 = "Where is the author's high school alma mater? Also, calculate 8 divided by 0, and remember my name."
    print(f"提问：{query3}")
    response3 = await agent.run(query3, ctx=ctx)
    print(f"回复：{response3}\n")

if __name__ == "__main__":
    asyncio.run(main())