# 客户端调用示例
import requests

# 文档上传
response = requests.post(
    "http://localhost:8080/insert",
    json={"input": "Your document content here"}
)

# 查询请求
response = requests.post(
    "http://localhost:8080/query",
    json={
        "query": "Your question",
        "mode": "hybrid",
        "stream": false
    }
)

# 支持多种存储后端
rag = LightRAG(
    vector_storage="MilvusStorage",        # Milvus 向量数据库
    graph_storage="Neo4jStorage",          # Neo4j 图数据库
    kv_storage="RedisStorage",             # Redis KV存储
    # 支持：FAISS, Qdrant, MongoDB, PostgreSQL 等
)

# 自定义 LLM 模型
from lightrag.llm.ollama import ollama_model_complete, ollama_embed

rag = LightRAG(
    llm_model_func=ollama_model_complete,
    embedding_func=ollama_embed,
    # 或使用 Azure OpenAI, Anthropic, Zhipu 等
)

# 核心分块算法 (lightrag/operate.py)
def chunking_by_token_size(
    tokenizer: Tokenizer,
    content: str,
    overlap_token_size: int = 128,
    max_token_size: int = 1024,
) -> list[dict]:
    tokens = tokenizer.encode(content)
    results = []
    for start in range(0, len(tokens), max_token_size - overlap_token_size):
        chunk_content = tokenizer.decode(tokens[start:start + max_token_size])
        results.append({
            "tokens": min(max_token_size, len(tokens) - start),
            "content": chunk_content.strip(),
            "chunk_order_index": index,
        })
    return results

# 实体关系抽取 Prompt 工程
ENTITY_EXTRACTION_PROMPT = """
从文本中提取实体和关系，输出格式：
("实体1", "实体类型", "实体描述")
("实体1", "关系类型", "实体2", "关系强度", "关系描述")
"""

async def naive_query(query, chunks_vdb, query_param, global_config):
    # 1. 查询向量化
    query_embedding = await embedding_func([query])
    
    # 2. 向量相似度检索
    results = await chunks_vdb.query(
        query, top_k=query_param.chunk_top_k
    )
    
    # 3. 重排序 (可选)
    if query_param.using_reranker:
        results = await rerank_func(query, results)
    
    return results

async def kg_query_local(query, entities_vdb, graph_storage):
    # 1. 检索相关实体
    similar_entities = await entities_vdb.query(query, top_k=top_k)
    
    # 2. 获取实体邻域关系
    related_relations = await graph_storage.get_nodes_edges_batch(
        [e["entity_name"] for e in similar_entities]
    )
    
    # 3. 加权融合实体+关系上下文
    context = merge_entity_relation_context(similar_entities, related_relations)
    return context

# 智能 Token 分配 (lightrag/operate.py)
def dynamic_context_allocation(
    max_total_tokens: int,
    kg_context_tokens: int,
    sys_prompt_overhead: int,
    buffer_tokens: int = 100
):
    available_chunk_tokens = max_total_tokens - kg_context_tokens - sys_prompt_overhead - buffer_tokens
    
    # 动态裁剪文本块，确保不超出模型上下文限制
    truncated_chunks = process_chunks_unified(
        chunks, chunk_token_limit=available_chunk_tokens
    )
    return truncated_chunks

# 大规模并发处理
async def batch_processing():
    tasks = [
        chunks_vdb.upsert(inserting_chunks),           # 向量存储
        process_entity_relation_graph(chunks),        # 图构建  
        full_docs.upsert(new_docs),                   # 文档存储
        text_chunks.upsert(inserting_chunks),         # 元数据存储
    ]
    await asyncio.gather(*tasks)  # 并行执行