# pip install openai
from openai import OpenAI

api_key = "xxx" # 如有密钥管理，用 api_key = os.getenv("DEEPSEEK_API_KEY)
base_url="https://xxx.xxx.com/v1/" # 注意 base_url 不要加 chat/completions

client = OpenAI(api_key=api_key, base_url=base_url)

# %% 简单调用
resp = client.chat.completions.create(
    model="DeepSeek-V3.2",
    messages=[
        {"role": "system", "content": "你是一个简洁的助手。"},
        {"role": "user", "content": "用一句话解释什么是YOLO标注格式。"},
    ],
)

print(resp.choices[0].message.content)

# %% 流式输出
stream = client.chat.completions.create(
    model="DeepSeek-V3.2",
    messages=[{"role": "user", "content": "写一个Python函数，判断素数。"}],
    stream=True,
)

for chunk in stream:
    delta = chunk.choices[0].delta
    if delta and delta.content:
        print(delta.content, end="", flush=True)
print()

# %% 多模态
import base64


def b64_image(path):
    return base64.b64encode(open(path, "rb").read()).decode()


resp = client.chat.completions.create(
    model="Qwen3-VL-32B-Instruct",
    messages=[{
        "role": "user",
        "content": [
            {"type": "text", "text": "这张图里有什么？"},
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64_image('1.png')}"}},
            # {"type": "image_url", "image_url": {"url": "https://xxx/1.png"}},
        ],
    }],
)

print(resp.choices[0].message.content)

#%% embedding

resp = client.embeddings.create(
    model="Qwen3-Embedding-0.6B",
    input=[
        "外卖不好吃",
        "送的很快",
        "太油了，太辣了"
    ],
)

# 每一条对应 resp.data 的一条
vec0 = resp.data[0].embedding  # list<float>

from mcp.server.fastmcp import FastMCP

mcp = FastMCP("get_info", log_level="ERROR")


@mcp.tool()
async def get_my_info(channel: str) -> str:
    # 下面的函数注释会先被 LLM 拿到
    """
    Get Guo Fei's public information by channel.

    Available channels (must match exactly):
    - "个人网站": personal website
    - "GitHub": GitHub profile
    - "知乎": Zhihu profile

    Args:
        channel: One of the predefined channel names above.

    Returns:
        A URL string corresponding to the requested channel.

    If an invalid channel is provided, a readable error message is returned.
    """

    info = {
        "个人网站": "https://www.guofei.site/"
        , "GitHub": "https://github.com/guofei9987/"
        , "知乎": "https://www.zhihu.com/people/guofei9987/answers/by_votes"
    }

    return info.get(channel, f"无效的渠道名称：{channel}。请使用以下之一：{', '.join(info.keys())}")


if __name__ == "__main__":
    mcp.run(transport='stdio')

query = "iPhone 15 Pro Max 多少钱"


payload = {
    'model': 'DeepSeek-V3.2',
    'messages': [
        {
            'role': 'system',
            'content': '你是一个助手'
        },
        {
            'role': 'user',
            'content': query
        }
    ],
}

res = call_model(payload)

print("未加入 RAG的结果：")
print(res)

# %% Chunking
def split_into_chunks(doc_file: str) -> List[str]:
    with open(doc_file, 'r') as f:
        content = f.read()
    return content.split('\n')


chunks = split_into_chunks("prices.txt")

# %% Embedding

from sentence_transformers import SentenceTransformer

embed_model = SentenceTransformer("./bge-small-zh-v1.5")


def embed_chunk(chunk: str) -> List[float]:
    embedding = embed_model.encode(chunk)
    return embedding.tolist()


test_embedding = embed_chunk("测试")
print("embedding 的维度", len(test_embedding))

embeddings = [embed_chunk(chunk) for chunk in chunks]

# meta 也可以参与召回
metas = [{"source": "prices.txt"} for chunk in chunks]

print("embedding 数量（片段个数）", len(embeddings))

# %% embedding 存入向量数据库中
import chromadb
from typing import List, Dict


# chromadb_client = chromadb.EphemeralClient() # 数据存在内存中
chromadb_client = chromadb.PersistentClient(path="./chroma_db")

chromadb_collection = chromadb_client.get_or_create_collection(name="default")


def save_embedding(chunks: List[str], embeddings: List[List[float]], metas: List[Dict]) -> None:
    ids = [str(i) for i in range(len(embeddings))]
    chromadb_collection.add(
        documents=chunks,
        embeddings=embeddings,
        ids=ids
        , metadatas=metas
    )


save_embedding(chunks, embeddings, metas)

from typing import List, Dict

from sentence_transformers import SentenceTransformer

embed_model = SentenceTransformer("./bge-small-zh-v1.5")


def embed_chunk(chunk: str) -> List[float]:
    embedding = embed_model.encode(chunk)
    return embedding.tolist()



#%% 加载向量数据库


import chromadb

chromadb_client = chromadb.PersistentClient(path="./chroma_db")

chromadb_collection = chromadb_client.get_or_create_collection(name="default")
print(f"数据库已加载 {chromadb_collection.count()} 条数据")

# %% 召回

def retrieve(query: str, top_k: int) -> List[str]:
    query_embedding = embed_chunk(query)
    results = chromadb_collection.query(
        query_embeddings=query_embedding,
        n_results=top_k
    )
    return results['documents'][0]


query = "iPhone 15 Pro Max 多少钱"

retrieved_chunks = retrieve(query, 10)

print("召回：", retrieved_chunks)

# %% 重排
from sentence_transformers import CrossEncoder

cross_encoder = CrossEncoder("./mmarco-mMiniLMv2-L12-H384-v1")


def rerank(query: str, retrieved_chunks: List[str], top_k: int = 3) -> List[str]:
    pairs = [(query, chunk) for chunk in retrieved_chunks]
    scores = cross_encoder.predict(pairs)

    chunk_with_scores = [(chunk, score)
                         for chunk, score in zip(retrieved_chunks, scores)]

    chunk_with_scores.sort(key=lambda x: x[1], reverse=True)

    return [chunk for chunk, _ in chunk_with_scores][:top_k]


reranked_chunks = rerank(query, retrieved_chunks)

print("重排后：", reranked_chunks)

# %% 测试结果：


query = "保温杯多少钱"
query = "皮鞋多少钱"
query = "How much is a milk"

retrieved_chunks = retrieve(query, top_k=10)
reranked_chunks = rerank(query, retrieved_chunks, top_k=3)

print("query:", query)
print("retrieved_chunks:", retrieved_chunks)
print("reranked_chunks:", reranked_chunks)


from openai import OpenAI

client = OpenAI(
    api_key="xxx",
    base_url="https://xxx.xxx.com/v1/"
)

resp = client.chat.completions.create(
    model="DeepSeek-V3.2",
    messages=[
        {
            'role': 'system',
            'content': f'''你是一个助手，请给予【相关材料】回答问题
# 【相关材料】
{"\n".join(reranked_chunks)}
            '''
        },
        {
            'role': 'user',
            'content': query
        }
    ],
)

print(resp.choices[0].message.content)

def save_bm25_model(chunks: List[str], model_path: str) -> None:
    tokenized_chunks = [tokenize_chinese(chunk) for chunk in chunks]
    bm25 = BM25Okapi(tokenized_chunks)

    with open(model_path, "wb") as f:
        pickle.dump({"bm25": bm25, "chunks": chunks}, f)

save_bm25_model(chunks, model_path='bm25.pkl')

def load_bm25_model(model_path: str) -> Tuple[List[str], BM25Okapi]:
    with open(model_path, 'rb') as f:
        data = pickle.load(f)

    bm25: BM25Okapi = data["bm25"]
    chunks: List[str] = data["chunks"]
    return chunks, bm25


# %%


def retrieve_bm25(query, bm25, chunks, k=5):
    tokenized_query = tokenize_chinese(query)
    bm25_scores = bm25.get_scores(tokenized_query)

    ranked_idx = sorted(range(len(chunks)), key=lambda i: bm25_scores[i], reverse=True)[:k]
    res = [(chunks[i], float(bm25_scores[i])) for i in ranked_idx]
    return res


chunks, bm25 = load_bm25_model("bm25.pkl")

k = 5
query = "iPhone 15 Pro Max 多少钱"
query = "羽绒服多少钱"

res = retrieve_bm25(query, bm25, chunks)

# 参考前面，传入 token
# client = OpenAI(...)


# 基于 prompt 生成文本
def get_completion(prompt, model="DeepSeek-V3.2"):
    resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}]
    )

    return resp.choices[0].message.content


# Query改写功能
class QueryRewriter:
    def __init__(self, model="DeepSeek-V3.2"):
        self.model = model

    def rewrite_context_dependent_query(self, current_query, conversation_history):
        """上下文依赖型Query改写"""
        instruction = """
你是一个智能的查询优化助手。请分析用户的当前问题以及前序对话历史，判断当前问题是否依赖于上下文。
如果依赖，请将当前问题改写成一个独立的、包含所有必要上下文信息的完整问题。
如果不依赖，直接返回原问题。
"""

        prompt = f"""
### 指令 ###
{instruction}

### 对话历史 ###
{conversation_history}

### 当前问题 ###
{current_query}

### 改写后的问题 ###
"""

        return get_completion(prompt, self.model)

    def rewrite_comparative_query(self, query, conversation_history):
        """对比型Query改写"""
        instruction = """
你是一个查询分析专家。请分析用户的输入和相关的对话上下文，识别出问题中需要进行比较的多个对象。
然后，将原始问题改写成一个更明确、更适合在知识库中检索的对比性查询。
"""

        prompt = f"""
### 指令 ###
{instruction}

### 对话历史 ###
{conversation_history}

### 原始问题 ###
{query}

### 改写后的查询 ###
"""

        return get_completion(prompt, self.model)

    def rewrite_ambiguous_reference_query(self, current_query, conversation_history):
        """模糊指代型Query改写"""
        instruction = """
你是一个消除语言歧义的专家。请分析用户的当前问题和对话历史，找出问题中 "都"、"它"、"这个" 等模糊指代词具体指向的对象。
然后，将这些指代词替换为明确的对象名称，生成一个清晰、无歧义的新问题。
"""

        prompt = f"""
### 指令 ###
{instruction}

### 对话历史 ###
{conversation_history}

### 当前问题 ###
{current_query}

### 改写后的问题 ###
"""

        return get_completion(prompt, self.model)

    def rewrite_multi_intent_query(self, current_query, conversation_history=None):
        """多意图型Query改写 - 分解查询
        ❗️ 多意图应当返回 List<str>，这个与别的不一样"""
        instruction = """
你是一个任务分解机器人。请将用户的复杂问题分解成多个独立的、可以单独回答的简单问题。以JSON数组格式输出。
"""

        prompt = f"""
### 指令 ###
{instruction}

### 原始问题 ###
{current_query}

### 分解后的问题列表 ###
请以JSON数组格式输出，例如：["问题1", "问题2", "问题3"]
"""

        return get_completion(prompt, self.model)

    def rewrite_rhetorical_query(self, current_query, conversation_history):
        """反问型Query改写"""
        instruction = """
你是一个沟通理解大师。请分析用户的反问或带有情绪的陈述，识别其背后真实的意图和问题。
然后，将这个反问改写成一个中立、客观、可以直接用于知识库检索的问题。
"""

        prompt = f"""
### 指令 ###
{instruction}

### 对话历史 ###
{conversation_history}

### 当前问题 ###
{current_query}

### 改写后的问题 ###
"""

        return get_completion(prompt, self.model)

    def rewrite(self, query, conversation_history, query_type: str):
        '''路由：根据 query_type 决定 调用'''
        if '上下文依赖' in query_type:
            final_result = self.rewrite_context_dependent_query(query, conversation_history)
        elif '对比' in query_type:
            final_result = self.rewrite_comparative_query(query, conversation_history)
        elif '模糊指代' in query_type:
            final_result = self.rewrite_ambiguous_reference_query(query, conversation_history)
        elif '多意图' in query_type:
            final_result = self.rewrite_multi_intent_query(query)
        elif '反问' in query_type:
            final_result = self.rewrite_rhetorical_query(query, conversation_history)
        else:
            final_result = query

        return final_result

    def query_type_classifier(self, query, conversation_history="", context_info=""):
        """自动识别Query类型并进行改写"""
        instruction = """
你是一个智能的查询分析专家。请分析用户的查询，识别其属于以下哪种类型：
1. 上下文依赖型 - 包含"还有"、"其他"等需要上下文理解的词汇
2. 对比型 - 包含"哪个"、"比较"、"更"、"哪个更好"、"哪个更"等比较词汇
3. 模糊指代型 - 包含"它"、"他们"、"都"、"这个"等指代词
4. 多意图型 - 包含多个独立问题，用"、"或"？"分隔
5. 反问型 - 包含"不会"、"难道"等反问语气
说明：如果同时存在多意图型、模糊指代型，优先级为多意图型>模糊指代型

请返回JSON格式的结果：
{
    "query_type": "查询类型",
    "confidence": "置信度(0-1)"
}
"""

        prompt = f"""
### 指令 ###
{instruction}

### 对话历史 ###
{conversation_history}

### 上下文信息 ###
{context_info}

### 原始查询 ###
{query}

### 分析结果 ###
"""

        response = get_completion(prompt, self.model)
        try:
            return json.loads(response)
        except:
            return {
                "query_type": "未知类型",
                "confidence": 0.5
            }

    def classify_and_rewrite(self, query, conversation_history="", context_info=""):
        query_type_predict = self.query_type_classifier(query, conversation_history, context_info)
        query_type = query_type_predict["query_type"]
        res = {"query_type": query_type_predict["query_type"]
            , "confidence": query_type_predict["confidence"]
            , "rewrite_result": self.rewrite(query, conversation_history, query_type)
               }
        return res


def main():
    # 初始化Query改写器
    rewriter = QueryRewriter()
    print("=== Query改写功能使用示例（迪士尼主题乐园） ===\n")
    cases = [
        {"sample_name": "示例1: 上下文依赖型Query"
            , "conversation_history": """
用户: "我想了解一下上海迪士尼乐园的最新项目。"
AI: "上海迪士尼乐园最新推出了'疯狂动物城'主题园区，这里有朱迪警官和尼克狐的互动体验。"
用户: "这个园区有什么游乐设施？"
AI: "'疯狂动物城'园区目前有疯狂动物城警察局、朱迪警官训练营和尼克狐的冰淇淋店等设施。"
"""
            , "current_query": "还有其他设施吗"
            , "query_type": "上下文依赖"
         }

        , {"sample_name": "示例2: 对比型Query"
            , "conversation_history": """
用户: "我想了解一下上海迪士尼乐园的最新项目。"
AI: "上海迪士尼乐园最新推出了疯狂动物城主题园区，还有蜘蛛侠主题园区"
"""
            , "current_query": "哪个游玩的时间比较长，比较有趣"
            , "query_type": "对比"

           }

        , {"sample_name": "示例3: 模糊指代型Query"
            , "conversation_history": """
用户: "我想了解一下上海迪士尼乐园和香港迪士尼乐园的烟花表演。"
AI: "好的，上海迪士尼乐园和香港迪士尼乐园都有精彩的烟花表演。"
"""
            , "current_query": "都什么时候开始？"
            , "query_type": "模糊指代"

           }

        , {"sample_name": "示例4: 多意图型Query"
            , "conversation_history": None
            , "current_query": "门票多少钱？需要提前预约吗？停车费怎么收？"
            , "query_type": "多意图"

           }
        , {"sample_name": "示例5: 反问型Query"
            , "conversation_history": """
用户: "你好，我想预订下周六上海迪士尼乐园的门票。"
AI: "正在为您查询... 查询到下周六的门票已经售罄。"
用户: "售罄是什么意思？我朋友上周去还能买到当天的票。"
"""
            , "current_query": "难道还要提前一个月预订吧？"
            , "query_type": "反问"

           }

    ]

    for case in cases:
        print('\n' + "=" * 20)
        conversation_history = case['conversation_history']
        current_query = case['current_query']
        query_type = case['query_type']
        print(case['sample_name'])
        print(f"对话历史: {conversation_history}")
        print(f"当前查询: {current_query}")

        type_predict = rewriter.query_type_classifier(current_query, conversation_history, query_type)
        print(f"预测的对话类型{type_predict}")
        result = rewriter.rewrite(current_query, conversation_history, query_type)
        print(f"改写结果: {result}\n")

        result = rewriter.classify_and_rewrite(current_query, conversation_history)
        print(f"自动分类+改写结果: {result}")


if __name__ == "__main__":
    main()

# 参考前面，传入 token
# client = OpenAI(...)


# 基于 prompt 生成文本
def get_completion(prompt, model="DeepSeek-V3.2"):
    resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}]
    )

    return resp.choices[0].message.content


# Query联网搜索改写功能
# 导入依赖库
import os
import json
import re
from datetime import datetime


class WebSearchQueryRewriter:
    def __init__(self, model="DeepSeek-V3.2"):
        self.model = model

    def identify_web_search_needs(self, query, conversation_history=""):
        """识别查询是否需要联网搜索"""
        instruction = """
你是一个智能的查询分析专家。请分析用户的查询，判断是否需要联网搜索来获取最新、最准确的信息。

需要联网搜索的情况包括：
1. 时效性信息 - 包含"最新"、"今天"、"现在"、"实时"、"当前"等时间相关词汇
2. 价格信息 - 包含"多少钱"、"价格"、"费用"、"票价"等价格相关词汇
3. 营业信息 - 包含"营业时间"、"开放时间"、"闭园时间"、"是否开放"等营业状态
4. 活动信息 - 包含"活动"、"表演"、"演出"、"节日"、"庆典"等动态信息
5. 天气信息 - 包含"天气"、"下雨"、"温度"等天气相关
6. 交通信息 - 包含"怎么去"、"交通"、"地铁"、"公交"等交通方式
7. 预订信息 - 包含"预订"、"预约"、"购票"、"订票"等预订相关
8. 实时状态 - 包含"排队"、"拥挤"、"人流量"等实时状态

请返回JSON格式：
{
    "need_web_search": true/false,
    "search_reason": "需要搜索的原因",
    "confidence": "置信度(0-1)"
}
"""

        prompt = f"""
### 指令 ###
{instruction}

### 对话历史 ###
{conversation_history}

### 用户查询 ###
{query}

### 分析结果 ###
"""

        response = get_completion(prompt, self.model)
        try:
            return json.loads(response)
        except:
            return {
                "need_web_search": False,
                "search_reason": "无法解析",
                "confidence": 0.5
            }

    def rewrite_for_web_search(self, query, search_type="general"):
        """为联网搜索改写查询"""
        instruction = """
你是一个专业的搜索查询优化专家。请将用户的查询改写为更适合搜索引擎检索的形式。

改写技巧：
1. 添加具体地点 - 如"上海迪士尼乐园"、"香港迪士尼乐园"
2. 添加时间范围 - 如"2024年"、"今天"、"本周"
3. 使用关键词组合 - 将长句拆分为关键词
4. 添加搜索意图 - 明确搜索目的
5. 去除口语化表达 - 转换为标准搜索词
6. 添加相关词汇 - 增加同义词或相关词

请返回JSON格式：
{
    "rewritten_query": "改写后的搜索查询",
    "search_keywords": ["关键词1", "关键词2", "关键词3"],
    "search_intent": "搜索意图",
    "suggested_sources": ["建议搜索的网站类型"]
}
"""

        prompt = f"""
### 指令 ###
{instruction}

### 原始查询 ###
{query}

### 搜索类型 ###
{search_type}

### 改写结果 ###
"""

        response = get_completion(prompt, self.model)
        try:
            return json.loads(response)
        except:
            return {
                "rewritten_query": query,
                "search_keywords": [query],
                "search_intent": "信息查询",
                "suggested_sources": ["官方网站", "旅游网站"]
            }

    def generate_search_strategy(self, query, search_type="general"):
        """生成搜索策略"""
        current_date = datetime.now().strftime("%Y年%m月%d日")
        instruction = f"""
你是一个搜索策略专家。请为用户的查询制定详细的搜索策略。

当前日期：{current_date}

搜索策略包括：
1. 主要搜索词 - 核心关键词
2. 扩展搜索词 - 相关词汇和同义词
3. 搜索网站 - 推荐的搜索平台
4. 时间范围 - 具体的搜索时间范围

请返回JSON格式：
primary_keywords
"""

        prompt = f"""
### 指令 ###
{instruction}

### 用户查询 ###
{query}

### 搜索类型 ###
{search_type}

### 搜索策略 ###
"""

        response = get_completion(prompt, self.model)
        try:
            return json.loads(response)
        except:
            return {
                "primary_keywords": [query],
                "extended_keywords": [],
                "search_platforms": ["百度", "谷歌"],
                "time_range": "最近一周"
            }

    def auto_web_search_rewrite(self, query, conversation_history=""):
        """自动识别并改写为联网搜索查询"""
        # 第一步：识别是否需要联网搜索
        search_analysis = self.identify_web_search_needs(query, conversation_history)

        if not search_analysis.get('need_web_search', False):
            return {
                "need_web_search": False,
                "reason": "查询不需要联网搜索",
                "original_query": query
            }

        # 第二步：改写查询
        rewritten_result = self.rewrite_for_web_search(query)

        # 第三步：生成搜索策略
        search_strategy = self.generate_search_strategy(query)

        return {
            "need_web_search": True,
            "search_reason": search_analysis.get('search_reason', ''),
            "confidence": search_analysis.get('confidence', 0.5),
            "original_query": query,
            "rewritten_query": rewritten_result.get('rewritten_query', query),
            "search_keywords": rewritten_result.get('search_keywords', []),
            "search_intent": rewritten_result.get('search_intent', ''),
            "suggested_sources": rewritten_result.get('suggested_sources', []),
            "search_strategy": search_strategy
        }


def main():
    # 初始化联网搜索Query改写器
    web_searcher = WebSearchQueryRewriter()

    print("=== Query联网搜索识别与改写示例（迪士尼主题乐园） ===\n")

    # 示例1: 时效性信息查询
    print("示例1: 时效性信息查询")
    conversation_history1 = """
用户: "我想去上海迪士尼乐园玩"
AI: "上海迪士尼乐园是一个很棒的选择！"
"""
    query1 = "上海迪士尼乐园今天开放吗？现在人多不多？"

    print(f"对话历史: {conversation_history1}")
    print(f"当前查询: {query1}")

    result1 = web_searcher.auto_web_search_rewrite(query1, conversation_history1)

    if result1['need_web_search']:
        print(f"✓ 需要联网搜索")
        print(f"  搜索原因: {result1['search_reason']}")
        print(f"  置信度: {result1['confidence']}")
        print(f"  改写查询: {result1['rewritten_query']}")
        print(f"  搜索关键词: {result1['search_keywords']}")
        print(f"  搜索意图: {result1['search_intent']}")
        print(f"  建议来源: {result1['suggested_sources']}")
        print(f"  搜索策略:")
        print(f"    - 主要关键词: {result1['search_strategy']['primary_keywords']}")
        print(f"    - 扩展关键词: {result1['search_strategy']['extended_keywords']}")
        print(f"    - 搜索平台: {result1['search_strategy']['search_platforms']}")
        print(f"    - 时间范围: {result1['search_strategy']['time_range']}")
    else:
        print(f"✗ 不需要联网搜索")
        print(f"  原因: {result1['reason']}")

    print("\n" + "=" * 60 + "\n")

    # 示例2: 价格和预订信息查询
    print("示例2: 价格和预订信息查询")
    query2 = "下周六的门票多少钱？需要提前多久预订？"

    print(f"当前查询: {query2}")

    result2 = web_searcher.auto_web_search_rewrite(query2)

    if result2['need_web_search']:
        print(f"✓ 需要联网搜索")
        print(f"  搜索原因: {result2['search_reason']}")
        print(f"  置信度: {result2['confidence']}")
        print(f"  改写查询: {result2['rewritten_query']}")
        print(f"  搜索关键词: {result2['search_keywords']}")
        print(f"  搜索意图: {result2['search_intent']}")
        print(f"  建议来源: {result2['suggested_sources']}")
        print(f"  搜索策略:")
        print(f"    - 主要关键词: {result2['search_strategy']['primary_keywords']}")
        print(f"    - 扩展关键词: {result2['search_strategy']['extended_keywords']}")
        print(f"    - 搜索平台: {result2['search_strategy']['search_platforms']}")
        print(f"    - 时间范围: {result2['search_strategy']['time_range']}")
    else:
        print(f"✗ 不需要联网搜索")
        print(f"  原因: {result2['reason']}")


if __name__ == "__main__":
    main()