Upload 4 files

Dockerfile

@@ -0,0 +1,25 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+RUN apt-get update && \
+    apt-get install -y git curl && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | bash && \
+    apt-get update && \
+    apt-get install -y git-lfs && \
+    git lfs install && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY . /app
+
+RUN pip install graphrag==1.0.0 fastapi uvicorn
+
+COPY search.py /usr/local/lib/python3.10/site-packages/graphrag/query/structured_search/local_search/search.py
+
+RUN chmod +x /app/start.sh
+
+EXPOSE 8080
+
+CMD ["/app/start.sh"]

search.py

@@ -0,0 +1,227 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+"""LocalSearch implementation."""
+
+import logging
+import time
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import tiktoken
+
+from graphrag.prompts.query.local_search_system_prompt import (
+    LOCAL_SEARCH_SYSTEM_PROMPT,
+)
+from graphrag.query.context_builder.builders import LocalContextBuilder
+from graphrag.query.context_builder.conversation_history import (
+    ConversationHistory,
+)
+from graphrag.query.llm.base import BaseLLM, BaseLLMCallback
+from graphrag.query.llm.text_utils import num_tokens
+from graphrag.query.structured_search.base import BaseSearch, SearchResult
+
+DEFAULT_LLM_PARAMS = {
+    "max_tokens": 1500,
+    "temperature": 0.0,
+}
+
+log = logging.getLogger(__name__)
+
+
+class LocalSearch(BaseSearch[LocalContextBuilder]):
+    """Search orchestration for local search mode."""
+
+    def __init__(
+        self,
+        llm: BaseLLM,
+        context_builder: LocalContextBuilder,
+        token_encoder: tiktoken.Encoding | None = None,
+        system_prompt: str | None = None,
+        response_type: str = "multiple paragraphs",
+        callbacks: list[BaseLLMCallback] | None = None,
+        llm_params: dict[str, Any] = DEFAULT_LLM_PARAMS,
+        context_builder_params: dict | None = None,
+    ):
+        super().__init__(
+            llm=llm,
+            context_builder=context_builder,
+            token_encoder=token_encoder,
+            llm_params=llm_params,
+            context_builder_params=context_builder_params or {},
+        )
+        self.system_prompt = system_prompt or LOCAL_SEARCH_SYSTEM_PROMPT
+        self.callbacks = callbacks
+        self.response_type = response_type
+
+    async def asearch(
+        self,
+        query: str,
+        conversation_history: ConversationHistory | None = None,
+        **kwargs,
+    ) -> SearchResult:
+        """Build local search context that fits a single context window and generate answer for the user query."""
+        start_time = time.time()
+        search_prompt = ""
+        llm_calls, prompt_tokens, output_tokens = {}, {}, {}
+        context_result = self.context_builder.build_context(
+            query=query,
+            conversation_history=conversation_history,
+            **kwargs,
+            **self.context_builder_params,
+        )
+        llm_calls["build_context"] = context_result.llm_calls
+        prompt_tokens["build_context"] = context_result.prompt_tokens
+        output_tokens["build_context"] = context_result.output_tokens
+
+        log.info("GENERATE ANSWER: %s. QUERY: %s", start_time, query)
+        try:
+            if "drift_query" in kwargs:
+                drift_query = kwargs["drift_query"]
+                search_prompt = self.system_prompt.format(
+                    context_data=context_result.context_chunks,
+                    response_type=self.response_type,
+                    global_query=drift_query,
+                )
+            else:
+                search_prompt = self.system_prompt.format(
+                    context_data=context_result.context_chunks,
+                    response_type=self.response_type,
+                )
+            search_messages = [
+                {"role": "system", "content": search_prompt},
+                {"role": "user", "content": query},
+            ]
+
+            response = await self.llm.agenerate(
+                messages=search_messages,
+                streaming=False,
+                callbacks=self.callbacks,
+                **self.llm_params,
+            )
+            llm_calls["response"] = 1
+            prompt_tokens["response"] = num_tokens(search_prompt, self.token_encoder)
+            output_tokens["response"] = num_tokens(response, self.token_encoder)
+
+            return SearchResult(
+                response=response,
+                context_data=context_result.context_records,
+                context_text=context_result.context_chunks,
+                completion_time=time.time() - start_time,
+                llm_calls=sum(llm_calls.values()),
+                prompt_tokens=sum(prompt_tokens.values()),
+                output_tokens=sum(output_tokens.values()),
+                llm_calls_categories=llm_calls,
+                prompt_tokens_categories=prompt_tokens,
+                output_tokens_categories=output_tokens,
+            )
+
+        except Exception:
+            log.exception("Exception in _asearch")
+            return SearchResult(
+                response="",
+                context_data=context_result.context_records,
+                context_text=context_result.context_chunks,
+                completion_time=time.time() - start_time,
+                llm_calls=1,
+                prompt_tokens=num_tokens(search_prompt, self.token_encoder),
+                output_tokens=0,
+            )
+
+    async def astream_search(
+        self,
+        query: str,
+        conversation_history: ConversationHistory | None = None,
+    ) -> AsyncGenerator:
+        """Build local search context that fits a single context window and generate answer for the user query."""
+        start_time = time.time()
+
+        context_result = self.context_builder.build_context(
+            query=query,
+            conversation_history=conversation_history,
+            **self.context_builder_params,
+        )
+        log.info("GENERATE ANSWER: %s. QUERY: %s", start_time, query)
+        search_prompt = self.system_prompt.format(
+            context_data=context_result.context_chunks, response_type=self.response_type
+        )
+        search_messages = [
+            {"role": "system", "content": search_prompt},
+            {"role": "user", "content": query},
+        ]
+
+        # send context records first before sending the reduce response
+        yield context_result.context_records
+        async for response in self.llm.astream_generate(  # type: ignore
+            messages=search_messages,
+            callbacks=self.callbacks,
+            **self.llm_params,
+        ):
+            yield response
+
+    def search(
+        self,
+        query: str,
+        conversation_history: ConversationHistory | None = None,
+        **kwargs,
+    ) -> SearchResult:
+        """Build local search context that fits a single context window and generate answer for the user question."""
+        start_time = time.time()
+        search_prompt = ""
+        llm_calls, prompt_tokens, output_tokens = {}, {}, {}
+        context_result = self.context_builder.build_context(
+            query=query,
+            conversation_history=conversation_history,
+            **kwargs,
+            **self.context_builder_params,
+        )
+        llm_calls["build_context"] = context_result.llm_calls
+        prompt_tokens["build_context"] = context_result.prompt_tokens
+        output_tokens["build_context"] = context_result.output_tokens
+
+        log.info("GENERATE ANSWER: %d. QUERY: %s", start_time, query)
+        try:
+            search_prompt = self.system_prompt.format(
+                context_data=context_result.context_chunks,
+                response_type=self.response_type,
+            )
+            search_messages = [
+                {"role": "system", "content": search_prompt},
+                {"role": "user", "content": query},
+            ]
+
+            response = self.llm.generate(
+                messages=search_messages,
+                streaming=True,
+                callbacks=self.callbacks,
+                **self.llm_params,
+            )
+            llm_calls["response"] = 1
+            prompt_tokens["response"] = num_tokens(search_prompt, self.token_encoder)
+            output_tokens["response"] = num_tokens(response, self.token_encoder)
+
+            return SearchResult(
+                response=response,
+                context_data=context_result.context_records,
+                context_text=context_result.context_chunks,
+                completion_time=time.time() - start_time,
+                llm_calls=sum(llm_calls.values()),
+                prompt_tokens=sum(prompt_tokens.values()),
+                output_tokens=sum(output_tokens.values()),
+                llm_calls_categories=llm_calls,
+                prompt_tokens_categories=prompt_tokens,
+                output_tokens_categories=output_tokens,
+            )
+
+        except Exception:
+            log.exception("Exception in _map_response_single_batch")
+            return SearchResult(
+                response="",
+                context_data=context_result.context_records,
+                context_text=context_result.context_chunks,
+                completion_time=time.time() - start_time,
+                llm_calls=1,
+                prompt_tokens=num_tokens(search_prompt, self.token_encoder),
+                output_tokens=0,
+            )
+

server.py

@@ -0,0 +1,378 @@
+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
+
+
+# 定义不同数据集的配置
+DATA_CONFIGS = {
+    "ghost": {
+        "input_dir": "/app/graphrag-data/data/the_bit_player",
+        "community_level": 2
+    },
+    "zhu_rongji": {
+        "input_dir": "/app/graphrag-data/data/the_bit_player", 
+        "community_level": 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)
+
+

start.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+cd /app
+git clone https://huggingface.co/datasets/nameliu/graphrag-data
+cd graphrag-data
+git checkout master
+git lfs pull
+
+cd /app
+python3 server.py