server.py

import os
import pandas as pd
import tiktoken

from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey
from graphrag.query.indexer_adapters import (
    read_indexer_covariates,
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.question_gen.local_gen import LocalQuestionGen
from graphrag.query.structured_search.local_search.mixed_context import (
    LocalSearchMixedContext,
)
from graphrag.query.structured_search.local_search.search import LocalSearch
from graphrag.vector_stores.lancedb import LanceDBVectorStore

# 基础数据目录路径
BASE_DATA_DIR = "/app/graphrag-data/data"

# 初始化DATA_CONFIGS字典
DATA_CONFIGS = {}

# 获取目录下所有文件夹
data_dirs = [d for d in os.listdir(BASE_DATA_DIR) if os.path.isdir(os.path.join(BASE_DATA_DIR, d))]

# 为每个文件夹创建配置
for dir_name in data_dirs:
    DATA_CONFIGS[dir_name] = {
        "input_dir": os.path.join(BASE_DATA_DIR, dir_name),
        "community_level": 2  # 默认值设置为2
    }

api_key = os.environ['api_key']
llm_model = os.environ['llm_model']
embedding_model = os.environ['embedding_model']
api_base = os.environ['api_base']

llm = ChatOpenAI(
    api_key=api_key,
    api_base=api_base,
    model=llm_model,
    api_type=OpenaiApiType.OpenAI,  # OpenaiApiType.OpenAI or OpenaiApiType.AzureOpenAI
    max_retries=10,
)

token_encoder = tiktoken.get_encoding("cl100k_base")

text_embedder = OpenAIEmbedding(
    api_key=api_key,
    api_base=api_base,
    api_type=OpenaiApiType.OpenAI,
    model=embedding_model,
    deployment_name=embedding_model,
    max_retries=7,
)

# 将数据加载逻辑封装成函数
def load_data(input_dir, community_level):
    lancedb_uri = f"{input_dir}/lancedb"
    
    # 定义表名
    COMMUNITY_REPORT_TABLE = "create_final_community_reports"
    ENTITY_TABLE = "create_final_nodes"
    ENTITY_EMBEDDING_TABLE = "create_final_entities"
    RELATIONSHIP_TABLE = "create_final_relationships"
    TEXT_UNIT_TABLE = "create_final_text_units"
    
    # 读取数据
    entity_df = pd.read_parquet(f"{input_dir}/{ENTITY_TABLE}.parquet")
    entity_embedding_df = pd.read_parquet(f"{input_dir}/{ENTITY_EMBEDDING_TABLE}.parquet")
    entities = read_indexer_entities(entity_df, entity_embedding_df, community_level)
    
    # 创建向量存储
    description_embedding_store = LanceDBVectorStore(
        collection_name="default-entity-description",
    )
    description_embedding_store.connect(db_uri=lancedb_uri)
    
    relationship_df = pd.read_parquet(f"{input_dir}/{RELATIONSHIP_TABLE}.parquet")
    relationships = read_indexer_relationships(relationship_df)
    
    report_df = pd.read_parquet(f"{input_dir}/{COMMUNITY_REPORT_TABLE}.parquet")
    reports = read_indexer_reports(report_df, entity_df, community_level)
    
    text_unit_df = pd.read_parquet(f"{input_dir}/{TEXT_UNIT_TABLE}.parquet")
    text_units = read_indexer_text_units(text_unit_df)
    
    return entities, description_embedding_store, relationships, reports, text_units

# 创建缓存字典来存储不同模型的搜索引擎实例
search_engines = {}

# 初始化函数
def initialize_search_engine(model_name):
    if model_name not in DATA_CONFIGS:
        raise ValueError(f"Unknown model: {model_name}")
        
    config = DATA_CONFIGS[model_name]
    # print(config)
    entities, description_embedding_store, relationships, reports, text_units = load_data(
        config["input_dir"], 
        config["community_level"]
    )
    
    context_builder = LocalSearchMixedContext(
        community_reports=reports,
        text_units=text_units,
        entities=entities,
        relationships=relationships,
        covariates=None,
        entity_text_embeddings=description_embedding_store,
        embedding_vectorstore_key=EntityVectorStoreKey.ID,
        text_embedder=text_embedder,
        token_encoder=token_encoder,
    )
    
    local_context_params = {
        "text_unit_prop": 0.5,
        "community_prop": 0.1,
        "conversation_history_max_turns": 5,
        "conversation_history_user_turns_only": True,
        "top_k_mapped_entities": 10,
        "top_k_relationships": 10,
        "include_entity_rank": True,
        "include_relationship_weight": True,
        "include_community_rank": False,
        "return_candidate_context": False,
        "embedding_vectorstore_key": EntityVectorStoreKey.ID,  # set this to EntityVectorStoreKey.TITLE if the vectorstore uses entity title as ids
        "max_tokens": 36_000,  # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)
    }
    
    llm_params = get_llm_params()
    return create_search_engine(llm, context_builder, token_encoder, llm_params, local_context_params)


from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import uvicorn
from datetime import datetime
import uuid
import time

app = FastAPI()

# 修改llm_params为动态配置
def get_llm_params(max_tokens=2000, temperature=0.0):
    return {
        "max_tokens": max_tokens,
        "temperature": temperature,
    }

def create_search_engine(llm, context_builder, token_encoder, llm_params, local_context_params):
    return LocalSearch(
        llm=llm,
        context_builder=context_builder,
        token_encoder=token_encoder,
        llm_params=llm_params,
        context_builder_params=local_context_params,
        response_type="multiple paragraphs",
    )


@app.post("/v1/completions")
async def completions(request: Request):
    body = await request.json()
    
    prompt = body.get("prompt", "hi")
    max_tokens = body.get("max_tokens", 2000)
    temperature = body.get("temperature", 0.0)
    model = body.get("model", "ghost")  # 默认使用ghost
    
    # 检查模型是否已初始化
    if model not in search_engines:
        try:
            search_engines[model] = initialize_search_engine(model)
        except ValueError as e:
            return JSONResponse(
                content={"error": str(e)},
                status_code=400
            )
    
    search_engine = search_engines[model]
    llm_params = get_llm_params(max_tokens, temperature)
    search_engine.llm_params = llm_params  # 更新LLM参数
    
    if prompt == "hi" or prompt == "":
        result_text = f"当前模型 {model} 已加载。可用模型: {', '.join(DATA_CONFIGS.keys())}"
        result = type('obj', (), {'response': result_text})()
    else:
        result = await search_engine.asearch(prompt)
    
    # 计算token使用情况（这里需要根据你的实际token计算方法进行修改）
    prompt_tokens = len(prompt.split())  # 简单示例，实际应使用proper tokenizer
    completion_tokens = len(result.response.split())
    total_tokens = prompt_tokens + completion_tokens

    # 构建响应
    response = {
        "id": f"cmpl-{str(uuid.uuid4())[:8]}",
        "object": "text_completion",
        "created": int(time.time()),
        "model": model,
        "system_fingerprint": f"fp_{str(uuid.uuid4())[:8]}",
        "choices": [
            {
                "text": result.response,
                "index": 0,
                "logprobs": None,
                "finish_reason": "length" if len(result.response.split()) >= max_tokens else "stop"
            }
        ],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens
        }
    }

    return JSONResponse(content=response)


from fastapi.responses import StreamingResponse
import json
import asyncio

@app.post("/api/v1/chat/completions")
async def chat_completions(request: Request):
    body = await request.json()

    # Extracting parameters from request body
    model = body.get("model", "ghost")  # Default model
    messages = body.get("messages", [])
    temperature = body.get("temperature", 0.0)
    max_tokens = body.get("max_tokens", 2000)
    stream = body.get("stream", False)  # 获取stream参数

    # Extracting user's prompt from messages
    user_message = next((msg["content"] for msg in messages if msg["role"] == "user"), "")

    # Check if the model exists in initialized search engines
    if model not in search_engines:
        try:
            search_engines[model] = initialize_search_engine(model)
        except ValueError as e:
            return JSONResponse(
                content={"error": str(e)},
                status_code=400
            )

    # Initialize search engine and LLM parameters
    search_engine = search_engines[model]
    llm_params = get_llm_params(max_tokens, temperature)
    search_engine.llm_params = llm_params

    # Handle 'empty' prompts to list available models
    if user_message == "" or user_message == "hi":
        result_text = f"当前模型 {model} 已加载。可用模型: {', '.join(DATA_CONFIGS.keys())}"
        result = type('obj', (), {'response': result_text})()
    else:
        # Fetch completions from search engine
        result = await search_engine.asearch(user_message)

    if not stream:
        # 非流式响应，返回完整的响应
        # Token usage calculation
        prompt_tokens = len(user_message.split())
        completion_tokens = len(result.response.split())
        total_tokens = prompt_tokens + completion_tokens

        completion_tokens_details = {
            "reasoning_tokens": 0,
            "accepted_prediction_tokens": 0,
            "rejected_prediction_tokens": 0
        }

        response = {
            "id": f"chatcmpl-{str(uuid.uuid4())[:8]}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": model,
            "usage": {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": total_tokens,
                "completion_tokens_details": completion_tokens_details
            },
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": result.response
                    },
                    "logprobs": None,
                    "finish_reason": "length" if len(result.response.split()) >= max_tokens else "stop",
                    "index": 0
                }
            ]
        }
        return JSONResponse(content=response)

    async def stream_response():
        chat_id = f"chatcmpl-{str(uuid.uuid4())[:8]}"
        system_fingerprint = f"fp_{str(uuid.uuid4())[:8]}"
        timestamp = int(time.time())

        # 发送role消息
        first_chunk = {
            'id': chat_id,
            'object': 'chat.completion.chunk',
            'created': timestamp,
            'model': model,
            'system_fingerprint': system_fingerprint,
            'choices': [{
                'index': 0,
                'delta': {'role': 'assistant'},
                'logprobs': None,
                'finish_reason': None
            }]
        }
        yield f"data: {json.dumps(first_chunk, ensure_ascii=False)}\n\n"

        # 将文本分成较大的块（每块约10个字符）
        text = result.response
        chunk_size = 50
        chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

        for chunk in chunks:
            data = {
                'id': chat_id,
                'object': 'chat.completion.chunk',
                'created': timestamp,
                'model': model,
                'system_fingerprint': system_fingerprint,
                'choices': [{
                    'index': 0,
                    'delta': {'content': chunk},
                    'logprobs': None,
                    'finish_reason': None
                }]
            }
            # 使用 ensure_ascii=False 确保中文正确显示
            json_str = json.dumps(data, ensure_ascii=False)
            yield f"data: {json_str}\n\n"
            await asyncio.sleep(0.1)  # 控制输出速度

        # 发送结束消息
        final_chunk = {
            'id': chat_id,
            'object': 'chat.completion.chunk',
            'created': timestamp,
            'model': model,
            'system_fingerprint': system_fingerprint,
            'choices': [{
                'index': 0,
                'delta': {},
                'logprobs': None,
                'finish_reason': 'stop'
            }]
        }
        yield f"data: {json.dumps(final_chunk, ensure_ascii=False)}\n\n"
        yield 'data: [DONE]\n\n'

    return StreamingResponse(
        stream_response(),
        media_type='text/event-stream'
    )

@app.get("/")
async def root():
    return "Hello from Docker!"

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080)