commit 7a596a0f636ae49f93ba2ca3272209a66a83b76a Author: xiaoxue Date: Sun May 24 23:37:54 2026 +0800 同步完整源码 - 2026-05-24 diff --git a/GraphRAG4OpenWebUI/test b/GraphRAG4OpenWebUI/test new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/GraphRAG4OpenWebUI/test @@ -0,0 +1 @@ + diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..4a878e3 --- /dev/null +++ b/README.md @@ -0,0 +1,161 @@ +### 🔥🔥🔥如有问题请联系我的微信 stoeng +### 🔥🔥🔥项目对应的视频演示请看 https://youtu.be/z4Si6O5NQ4c + +# GraphRAG4OpenWebUI +
+

Integrate Microsoft's GraphRAG Technology into Open WebUI for Advanced Information Retrieval

+ English | 简体中文 +
+ +GraphRAG4OpenWebUI is an API interface specifically designed for Open WebUI, aiming to integrate Microsoft Research's GraphRAG (Graph-based Retrieval-Augmented Generation) technology. This project provides a powerful information retrieval system that supports multiple search models, particularly suitable for use in open web user interfaces. + +## Project Overview + +The main goal of this project is to provide a convenient interface for Open WebUI to leverage the powerful features of GraphRAG. It integrates three main retrieval methods and offers a comprehensive search option, allowing users to obtain thorough and precise search results. + +### Key Retrieval Features + +1. **Local Search** + - Utilizes GraphRAG technology for efficient retrieval in local knowledge bases + - Suitable for quick access to pre-defined structured information + - Leverages graph structures to improve retrieval accuracy and relevance + +2. **Global Search** + - Searches for information in a broader scope, beyond local knowledge bases + - Suitable for queries requiring more comprehensive information + - Utilizes GraphRAG's global context understanding capabilities to provide richer search results + +3. **Tavily Search** + - Integrates external Tavily search API + - Provides additional internet search capabilities, expanding information sources + - Suitable for queries requiring the latest or extensive web information + +4. **Full Model Search** + - Combines all three search methods above + - Provides the most comprehensive search results, meeting complex information needs + - Automatically integrates and ranks information from different sources + +### Local LLM and Embedding Model Support + +GraphRAG4OpenWebUI now supports the use of local language models (LLMs) and embedding models, increasing the project's flexibility and privacy. Specifically, we support the following local models: + +1. **Ollama** + - Supports various open-source LLMs run through Ollama, such as Llama 2, Mistral, etc. + - Can be configured by setting the `API_BASE` environment variable to point to Ollama's API endpoint + +2. **LM Studio** + - Compatible with models run by LM Studio + - Connect to LM Studio's service by configuring the `API_BASE` environment variable + +3. **Local Embedding Models** + - Supports the use of locally run embedding models, such as SentenceTransformers + - Specify the embedding model to use by setting the `GRAPHRAG_EMBEDDING_MODEL` environment variable + +This support for local models allows GraphRAG4OpenWebUI to run without relying on external APIs, enhancing data privacy and reducing usage costs. + +## Installation +Ensure that you have Python 3.8 or higher installed on your system. Then, follow these steps to install: +1. Clone the repository: + ```bash + git clone https://github.com/your-username/GraphRAG4OpenWebUI.git + cd GraphRAG4OpenWebUI + ``` + +2. Create and activate a virtual environment: + ```bash + python -m venv venv + source venv/bin/activate # On Windows use venv\Scripts\activate + ``` + +3. Install dependencies: + ```bash + pip install -r requirements.txt + ``` + Note: The graphrag package might need to be installed from a specific source. If the above command fails to install graphrag, please refer to Microsoft Research's specific instructions or contact the maintainer for the correct installation method. + +## Configuration + +Before running the API, you need to set the following environment variables. You can do this by creating a `.env` file or exporting them directly in your terminal: + + + +```bash +# Set the TAVILY API key +export TAVILY_API_KEY="your_tavily_api_key" + +export INPUT_DIR="/path/to/your/input/directory" + +# Set the API key for LLM +export GRAPHRAG_API_KEY="your_actual_api_key_here" + +# Set the API key for embedding (if different from GRAPHRAG_API_KEY) +export GRAPHRAG_API_KEY_EMBEDDING="your_embedding_api_key_here" + +# Set the LLM model +export GRAPHRAG_LLM_MODEL="gemma2" + +# Set the API base URL +export API_BASE="http://localhost:11434/v1" + +# Set the embedding API base URL (default is OpenAI's API) +export API_BASE_EMBEDDING="https://api.openai.com/v1" + +# Set the embedding model (default is "text-embedding-3-small") +export GRAPHRAG_EMBEDDING_MODEL="text-embedding-3-small" +``` + +Make sure to replace the placeholders in the above commands with your actual API keys and paths. + +## Usage + +1. Start the server: + ``` + python main-en.py + ``` + The server will run on `http://localhost:8012`. + +2. API Endpoints: + - `/v1/chat/completions`: POST request for performing searches + - `/v1/models`: GET request to retrieve the list of available models + +3. Integration with Open WebUI: + In the Open WebUI configuration, set the API endpoint to `http://localhost:8012/v1/chat/completions`. This will allow Open WebUI to use the search functionality of GraphRAG4OpenWebUI. + +4. Example search request: + ```python + import requests + import json + + url = "http://localhost:8012/v1/chat/completions" + headers = {"Content-Type": "application/json"} + data = { + "model": "full-model:latest", + "messages": [{"role": "user", "content": "Your search query"}], + "temperature": 0.7 + } + + response = requests.post(url, headers=headers, data=json.dumps(data)) + print(response.json()) + ``` + +## Available Models + +- `graphrag-local-search:latest`: Local search +- `graphrag-global-search:latest`: Global search +- `tavily-search:latest`: Tavily search +- `full-model:latest`: Comprehensive search (includes all search methods above) + +## Notes + +- Ensure that you have the correct input files (such as Parquet files) in the `INPUT_DIR` directory. +- The API uses asynchronous programming, make sure your environment supports async operations. +- For large-scale deployment, consider using a production-grade ASGI server. +- This project is specifically designed for Open WebUI and can be easily integrated into various web-based applications. + +## Contributing + +Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change. + +## License + +[Apache-2.0 License](LICENSE) diff --git a/README_ZH-CN.md b/README_ZH-CN.md new file mode 100644 index 0000000..3799451 --- /dev/null +++ b/README_ZH-CN.md @@ -0,0 +1,159 @@ +### 🔥🔥🔥如有问题请联系我的微信 stoeng +### 🔥🔥🔥项目对应的视频演示请看 https://youtu.be/z4Si6O5NQ4c + +# GraphRAG4OpenWebUI +
+

将微软的 GraphRAG 技术集成到 Open WebUI 中，实现高级信息检索

+ English | 简体中文 +
+ +GraphRAG4OpenWebUI 是一个专为 Open WebUI 设计的 API 接口，旨在集成微软研究院的 GraphRAG（基于图的检索增强生成）技术。该项目提供了一个强大的信息检索系统，支持多种搜索模型，特别适合在开放式 Web 用户界面中使用。 + +## 项目概述 + +本项目的主要目标是为 Open WebUI 提供一个便捷的接口，以利用 GraphRAG 的强大功能。它集成了三种主要的检索方法，并提供了一个综合搜索选项，使用户能够获得全面而精确的搜索结果。 + +### 主要检索功能 + +1. **本地搜索（Local Search）** + - 利用 GraphRAG 技术在本地知识库中进行高效检索 + - 适用于快速访问预先定义的结构化信息 + - 利用图结构提高检索的准确性和相关性 + +2. **全局搜索（Global Search）** + - 在更广泛的范围内搜索信息，超越本地知识库的限制 + - 适用于需要更全面信息的查询 + - 利用 GraphRAG 的全局上下文理解能力，提供更丰富的搜索结果 + +3. **Tavily 搜索** + - 集成外部 Tavily 搜索 API + - 提供额外的互联网搜索能力，扩展信息源 + - 适用于需要最新或广泛网络信息的查询 + +4. **全模型搜索（Full Model Search）** + - 综合上述三种搜索方法 + - 提供最全面的搜索结果，满足复杂的信息需求 + - 自动整合和排序来自不同来源的信息 + +### 本地LLM和Embedding模型支持 + +GraphRAG4OpenWebUI 现在支持使用本地的语言模型（LLM）和嵌入模型，增加了项目的灵活性和隐私性。特别地，我们支持以下本地模型： + +1. **Ollama** + - 支持使用 Ollama 运行的各种开源 LLM，如 Llama 2、Mistral 等 + - 可以通过设置 `API_BASE` 环境变量来指向 Ollama 的 API 端点 + +2. **LM Studio** + - 兼容 LM Studio 运行的模型 + - 通过配置 `API_BASE` 环境变量连接到 LM Studio 的服务 + +3. **本地 Embedding 模型** + - 支持使用本地运行的嵌入模型，如 SentenceTransformers + - 通过设置 `GRAPHRAG_EMBEDDING_MODEL` 环境变量来指定使用的嵌入模型 + +这些本地模型的支持使得 GraphRAG4OpenWebUI 能够在不依赖外部API的情况下运行，提高了数据隐私和降低了使用成本。 + +## 安装 +确保您的系统中已安装 Python 3.8 或更高版本。然后，按照以下步骤安装： +1. 克隆仓库： + ```bash + git clone https://github.com/your-username/GraphRAG4OpenWebUI.git + cd GraphRAG4OpenWebUI + ``` + +2. 创建并激活虚拟环境： + ```bash + python -m venv venv + source venv/bin/activate # 在 Windows 上使用 venv\Scripts\activate + ``` + +3. 安装依赖： + ```bash + pip install -r requirements.txt + ``` + 注意：graphrag 包可能需要从特定的源安装。如果上述命令无法安装 graphrag，请参考微软研究院的具体说明或联系维护者获取正确的安装方法。 + +## 配置 + +在运行 API 之前，需要设置以下环境变量。您可以通过创建 `.env` 文件或直接在终端中导出这些变量： + + +```bash +export TAVILY_API_KEY="your_tavily_api_key" + +export INPUT_DIR="/path/to/your/input/directory" + +# 设置llm API密钥 +export GRAPHRAG_API_KEY="your_actual_api_key_here" + +# 设置嵌入API密钥（如果与GRAPHRAG_API_KEY不同） +export GRAPHRAG_API_KEY_EMBEDDING="your_embedding_api_key_here" + +# 设置LLM模型（默认为"gemma2"） +export GRAPHRAG_LLM_MODEL="gemma2" + +# 设置API基础URL（默认为本地服务器） +export API_BASE="http://localhost:11434/v1" + +# 设置嵌入API基础URL（默认为OpenAI的API） +export API_BASE_EMBEDDING="https://api.openai.com/v1" + +# 设置嵌入模型（默认为"text-embedding-3-small"） +export GRAPHRAG_EMBEDDING_MODEL="text-embedding-3-small" +``` + +请确保将上述命令中的占位符替换为实际的 API 密钥和路径。 + +## 使用方法 + +1. 启动服务器： + ``` + python main-cn.py + ``` + 服务器将在 `http://localhost:8012` 上运行。 + +2. API 端点： + - `/v1/chat/completions`: POST 请求，用于执行搜索 + - `/v1/models`: GET 请求，获取可用模型列表 + +3. 在 Open WebUI 中集成： + 在 Open WebUI 的配置中，将 API 端点设置为 `http://localhost:8012/v1/chat/completions`。这将允许 Open WebUI 使用 GraphRAG4OpenWebUI 的搜索功能。 + +4. 发送搜索请求示例： + ```python + import requests + import json + + url = "http://localhost:8012/v1/chat/completions" + headers = {"Content-Type": "application/json"} + data = { + "model": "full-model:latest", + "messages": [{"role": "user", "content": "您的搜索查询"}], + "temperature": 0.7 + } + + response = requests.post(url, headers=headers, data=json.dumps(data)) + print(response.json()) + ``` + +## 可用模型 + +- `graphrag-local-search:latest`: 本地搜索 +- `graphrag-global-search:latest`: 全局搜索 +- `tavily-search:latest`: Tavily 搜索 +- `full-model:latest`: 综合搜索（包含上述所有搜索方法） + +## 注意事项 + +- 确保在 `INPUT_DIR` 目录中有正确的输入文件（如 Parquet 文件）。 +- API 使用异步编程，确保您的环境支持异步操作。 +- 对于大规模部署，建议使用生产级的 ASGI 服务器。 +- 本项目专为 Open WebUI 设计，可以轻松集成到各种基于 Web 的应用中。 + +## 贡献 + +我们欢迎您提交 Pull Requests 来改进这个项目。对于重大变更，请先开 issue 讨论您想要改变的内容。 + +## 许可证 + +[Apache-2.0 许可证](LICENSE) diff --git a/graphrag3dknowledge.py b/graphrag3dknowledge.py new file mode 100644 index 0000000..0146663 --- /dev/null +++ b/graphrag3dknowledge.py @@ -0,0 +1,293 @@ +import os #用于文件系统操作 +import pandas as pd #用于数据处理和操作 +import networkx as nx #用于创建和分析图结构 +import plotly.graph_objects as go #plotly：用于创建交互式可视化 plotly.graph_objects：用于创建低级的plotly图形对象 +from plotly.subplots import make_subplots #用于创建子图 +import plotly.express as px #用于快速创建统计图表 + +def read_parquet_files(directory): + """ + 读取指定目录下的所有Parquet文件并合并 + 功能：读取指定目录下的所有Parquet文件并合并成一个DataFrame + 实现：使用os.listdir遍历目录，pd.read_parquet读取每个文件，然后用pd.concat合并 + """ + dataframes = [] + for filename in os.listdir(directory): + if filename.endswith('.parquet'): + file_path = os.path.join(directory, filename) + df = pd.read_parquet(file_path) + dataframes.append(df) + return pd.concat(dataframes, ignore_index=True) if dataframes else pd.DataFrame() + + +def clean_dataframe(df): + """ + 清理DataFrame，移除无效的行 + 功能：清理DataFrame，移除无效的行 + 实现：删除source和target列中的空值，将这两列转换为字符串类型 + """ + df = df.dropna(subset=['source', 'target']) + df['source'] = df['source'].astype(str) + df['target'] = df['target'].astype(str) + return df + + +def create_knowledge_graph(df): + """ + 从DataFrame创建知识图谱 + 功能：从DataFrame创建知识图谱 + 实现：使用networkx创建有向图，遍历DataFrame的每一行，添加边和属性 + """ + G = nx.DiGraph() + for _, row in df.iterrows(): + source = row['source'] + target = row['target'] + attributes = {k: v for k, v in row.items() if k not in ['source', 'target']} + G.add_edge(source, target, **attributes) + return G + + +def create_node_link_trace(G, pos): + """ + 功能：创建节点和边的3D轨迹 + 实现：使用networkx的布局信息创建Plotly的Scatter3d对象 + """ + edge_x = [] + edge_y = [] + edge_z = [] + for edge in G.edges(): + x0, y0, z0 = pos[edge[0]] + x1, y1, z1 = pos[edge[1]] + edge_x.extend([x0, x1, None]) + edge_y.extend([y0, y1, None]) + edge_z.extend([z0, z1, None]) + + edge_trace = go.Scatter3d( + x=edge_x, y=edge_y, z=edge_z, + line=dict(width=0.5, color='#888'), + hoverinfo='none', + mode='lines') + + node_x = [pos[node][0] for node in G.nodes()] + node_y = [pos[node][1] for node in G.nodes()] + node_z = [pos[node][2] for node in G.nodes()] + + node_trace = go.Scatter3d( + x=node_x, y=node_y, z=node_z, + mode='markers', + hoverinfo='text', + marker=dict( + showscale=True, + colorscale='YlGnBu', + size=10, + colorbar=dict( + thickness=15, + title='Node Connections', + xanchor='left', + titleside='right' + ) + ) + ) + + node_adjacencies = [] + node_text = [] + for node, adjacencies in G.adjacency(): + node_adjacencies.append(len(adjacencies)) + node_text.append(f'Node: {node}
# of connections: {len(adjacencies)}') + + node_trace.marker.color = node_adjacencies + node_trace.text = node_text + + return edge_trace, node_trace + + +def create_edge_label_trace(G, pos, edge_labels): + """ + 功能：创建边标签的3D轨迹 + 实现：计算边的中点位置，创建Scatter3d对象显示标签 + """ + return go.Scatter3d( + x=[pos[edge[0]][0] + (pos[edge[1]][0] - pos[edge[0]][0]) / 2 for edge in edge_labels], + y=[pos[edge[0]][1] + (pos[edge[1]][1] - pos[edge[0]][1]) / 2 for edge in edge_labels], + z=[pos[edge[0]][2] + (pos[edge[1]][2] - pos[edge[0]][2]) / 2 for edge in edge_labels], + mode='text', + text=list(edge_labels.values()), + textposition='middle center', + hoverinfo='none' + ) + + +def create_degree_distribution(G): + """ + 功能：创建节点度分布直方图 + 实现：使用plotly.express创建直方图 + """ + degrees = [d for n, d in G.degree()] + fig = px.histogram(x=degrees, nbins=20, labels={'x': 'Degree', 'y': 'Count'}) + fig.update_layout( + title_text='Node Degree Distribution', + margin=dict(l=0, r=0, t=30, b=0), + height=300 + ) + return fig + + +def create_centrality_plot(G): + """ + 功能：创建节点中心性分布箱线图 + 实现：计算度中心性，使用plotly.express创建箱线图 + """ + centrality = nx.degree_centrality(G) + centrality_values = list(centrality.values()) + fig = px.box(y=centrality_values, labels={'y': 'Centrality'}) + fig.update_layout( + title_text='Degree Centrality Distribution', + margin=dict(l=0, r=0, t=30, b=0), + height=300 + ) + return fig + + +def visualize_graph_plotly(G): + """功能：使用Plotly创建全面优化布局的高级交互式知识图谱可视化 + 实现： + 创建3D布局 + 生成节点和边的轨迹 + 创建子图，包括3D图、度分布图和中心性分布图 + 添加交互式按钮和滑块 + 优化整体布局 + """ + if G.number_of_nodes() == 0: + print("Graph is empty. Nothing to visualize.") + return + + pos = nx.spring_layout(G, dim=3) # 3D layout + edge_trace, node_trace = create_node_link_trace(G, pos) + + edge_labels = nx.get_edge_attributes(G, 'relation') + edge_label_trace = create_edge_label_trace(G, pos, edge_labels) + + degree_dist_fig = create_degree_distribution(G) + centrality_fig = create_centrality_plot(G) + + fig = make_subplots( + rows=2, cols=2, + column_widths=[0.7, 0.3], + row_heights=[0.7, 0.3], + specs=[ + [{"type": "scene", "rowspan": 2}, {"type": "xy"}], + [None, {"type": "xy"}] + ], + subplot_titles=("3D Knowledge Graph Code by AI超元域频道", "Node Degree Distribution", "Degree Centrality Distribution") + ) + + fig.add_trace(edge_trace, row=1, col=1) + fig.add_trace(node_trace, row=1, col=1) + fig.add_trace(edge_label_trace, row=1, col=1) + + fig.add_trace(degree_dist_fig.data[0], row=1, col=2) + fig.add_trace(centrality_fig.data[0], row=2, col=2) + + # Update 3D layout + fig.update_layout( + scene=dict( + xaxis=dict(showticklabels=False, showgrid=False, zeroline=False), + yaxis=dict(showticklabels=False, showgrid=False, zeroline=False), + zaxis=dict(showticklabels=False, showgrid=False, zeroline=False), + aspectmode='cube' + ), + scene_camera=dict(eye=dict(x=1.5, y=1.5, z=1.5)) + ) + + # Add buttons for different layouts + fig.update_layout( + updatemenus=[ + dict( + type="buttons", + direction="left", + buttons=list([ + dict(args=[{"visible": [True, True, True, True, True]}], label="Show All", method="update"), + dict(args=[{"visible": [True, True, False, True, True]}], label="Hide Edge Labels", + method="update"), + dict(args=[{"visible": [False, True, False, True, True]}], label="Nodes Only", method="update") + ]), + pad={"r": 10, "t": 10}, + showactive=True, + x=0.05, + xanchor="left", + y=1.1, + yanchor="top" + ), + ] + ) + + # Add slider for node size + fig.update_layout( + sliders=[dict( + active=0, + currentvalue={"prefix": "Node Size: "}, + pad={"t": 50}, + steps=[dict(method='update', + args=[{'marker.size': [i] * len(G.nodes)}], + label=str(i)) for i in range(5, 21, 5)] + )] + ) + + # 优化整体布局 + # fig.update_layout( + # height=1198, # 增加整体高度 + # width=2055, # 增加整体宽度 + # title_text="Advanced Interactive Knowledge Graph", + # margin=dict(l=10, r=10, t=25, b=10), + # legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01) + # ) + + fig.show() + + +def main(): + """ 功能：主函数，协调整个程序的执行流程 + 实现： + 读取Parquet文件 + 清理数据 + 创建知识图谱 + 打印图的统计信息 + 调用可视化函数 + """ + directory = '/Users/charlesqin/PycharmProjects/RAGCode/inputs/artifacts' # 替换为实际的目录路径 + df = read_parquet_files(directory) + + if df.empty: + print("No data found in the specified directory.") + return + + print("Original DataFrame shape:", df.shape) + print("Original DataFrame columns:", df.columns.tolist()) + print("Original DataFrame head:") + print(df.head()) + + df = clean_dataframe(df) + + print("\nCleaned DataFrame shape:", df.shape) + print("Cleaned DataFrame head:") + print(df.head()) + + if df.empty: + print("No valid data remaining after cleaning.") + return + + G = create_knowledge_graph(df) + + print(f"\nGraph statistics:") + print(f"Nodes: {G.number_of_nodes()}") + print(f"Edges: {G.number_of_edges()}") + + if G.number_of_nodes() > 0: + print(f"Connected components: {nx.number_connected_components(G.to_undirected())}") + visualize_graph_plotly(G) + else: + print("Graph is empty. Cannot visualize.") + + +if __name__ == "__main__": + main() diff --git a/main-cn.py b/main-cn.py new file mode 100644 index 0000000..69ed1f0 --- /dev/null +++ b/main-cn.py @@ -0,0 +1,503 @@ +import os +import asyncio +import time +import uuid +import json +import re +import pandas as pd +import tiktoken +import logging +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse +from pydantic import BaseModel, Field +from typing import List, Optional, Dict, Any, Union +from contextlib import asynccontextmanager +from tavily import TavilyClient + + +# GraphRAG 相关导入 +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.input.loaders.dfs import store_entity_semantic_embeddings +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.query.structured_search.global_search.community_context import GlobalCommunityContext +from graphrag.query.structured_search.global_search.search import GlobalSearch +from graphrag.vector_stores.lancedb import LanceDBVectorStore + +# 设置日志 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + +# 设置常量和配置 +INPUT_DIR = os.getenv('INPUT_DIR') +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" +COVARIATE_TABLE = "create_final_covariates" +TEXT_UNIT_TABLE = "create_final_text_units" +COMMUNITY_LEVEL = 2 +PORT = 8012 + +# 全局变量，用于存储搜索引擎和问题生成器 +local_search_engine = None +global_search_engine = None +question_generator = None + + +# 数据模型 +class Message(BaseModel): + role: str + content: str + + +class ChatCompletionRequest(BaseModel): + model: str + messages: List[Message] + temperature: Optional[float] = 1.0 + top_p: Optional[float] = 1.0 + n: Optional[int] = 1 + stream: Optional[bool] = False + stop: Optional[Union[str, List[str]]] = None + max_tokens: Optional[int] = None + presence_penalty: Optional[float] = 0 + frequency_penalty: Optional[float] = 0 + logit_bias: Optional[Dict[str, float]] = None + user: Optional[str] = None + + +class ChatCompletionResponseChoice(BaseModel): + index: int + message: Message + finish_reason: Optional[str] = None + + +class Usage(BaseModel): + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletionResponse(BaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}") + object: str = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseChoice] + usage: Usage + system_fingerprint: Optional[str] = None + + +async def setup_llm_and_embedder(): + """ + 设置语言模型（LLM）和嵌入模型 + """ + logger.info("正在设置LLM和嵌入器") + + # 获取API密钥和基础URL + api_key = os.environ.get("GRAPHRAG_API_KEY", "YOUR_API_KEY") + api_key_embedding = os.environ.get("GRAPHRAG_API_KEY_EMBEDDING", api_key) + api_base = os.environ.get("API_BASE", "https://api.openai.com/v1") + api_base_embedding = os.environ.get("API_BASE_EMBEDDING", "https://api.openai.com/v1") + + # 获取模型名称 + llm_model = os.environ.get("GRAPHRAG_LLM_MODEL", "gpt-3.5-turbo-0125") + embedding_model = os.environ.get("GRAPHRAG_EMBEDDING_MODEL", "text-embedding-3-small") + + # 检查API密钥是否存在 + if api_key == "YOUR_API_KEY": + logger.error("环境变量中未找到有效的GRAPHRAG_API_KEY") + raise ValueError("GRAPHRAG_API_KEY未正确设置") + + # 初始化ChatOpenAI实例 + llm = ChatOpenAI( + api_key=api_key, + api_base=api_base, + model=llm_model, + api_type=OpenaiApiType.OpenAI, + max_retries=20, + ) + + # 初始化token编码器 + token_encoder = tiktoken.get_encoding("cl100k_base") + + # 初始化文本嵌入模型 + text_embedder = OpenAIEmbedding( + api_key=api_key_embedding, + api_base=api_base_embedding, + api_type=OpenaiApiType.OpenAI, + model=embedding_model, + deployment_name=embedding_model, + max_retries=20, + ) + + + logger.info("LLM和嵌入器设置完成") + return llm, token_encoder, text_embedder + + +async def load_context(): + """ + 加载上下文数据，包括实体、关系、报告、文本单元和协变量 + """ + logger.info("正在加载上下文数据") + try: + 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="entity_description_embeddings") + description_embedding_store.connect(db_uri=LANCEDB_URI) + store_entity_semantic_embeddings(entities=entities, vectorstore=description_embedding_store) + + 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) + + covariate_df = pd.read_parquet(f"{INPUT_DIR}/{COVARIATE_TABLE}.parquet") + claims = read_indexer_covariates(covariate_df) + logger.info(f"声明记录数: {len(claims)}") + covariates = {"claims": claims} + + logger.info("上下文数据加载完成") + return entities, relationships, reports, text_units, description_embedding_store, covariates + except Exception as e: + logger.error(f"加载上下文数据时出错: {str(e)}") + raise + + +async def setup_search_engines(llm, token_encoder, text_embedder, entities, relationships, reports, text_units, + description_embedding_store, covariates): + """ + 设置本地搜索引擎和全局搜索引擎 + """ + logger.info("正在设置搜索引擎") + + # 设置本地搜索引擎 + local_context_builder = LocalSearchMixedContext( + community_reports=reports, + text_units=text_units, + entities=entities, + relationships=relationships, + covariates=covariates, + 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, + "max_tokens": 12_000, + } + + local_llm_params = { + "max_tokens": 2_000, + "temperature": 0.0, + } + + local_search_engine = LocalSearch( + llm=llm, + context_builder=local_context_builder, + token_encoder=token_encoder, + llm_params=local_llm_params, + context_builder_params=local_context_params, + response_type="multiple paragraphs", + ) + + # 设置全局搜索引擎 + global_context_builder = GlobalCommunityContext( + community_reports=reports, + entities=entities, + token_encoder=token_encoder, + ) + + global_context_builder_params = { + "use_community_summary": False, + "shuffle_data": True, + "include_community_rank": True, + "min_community_rank": 0, + "community_rank_name": "rank", + "include_community_weight": True, + "community_weight_name": "occurrence weight", + "normalize_community_weight": True, + "max_tokens": 12_000, + "context_name": "Reports", + } + + map_llm_params = { + "max_tokens": 1000, + "temperature": 0.0, + "response_format": {"type": "json_object"}, + } + + reduce_llm_params = { + "max_tokens": 2000, + "temperature": 0.0, + } + + global_search_engine = GlobalSearch( + llm=llm, + context_builder=global_context_builder, + token_encoder=token_encoder, + max_data_tokens=12_000, + map_llm_params=map_llm_params, + reduce_llm_params=reduce_llm_params, + allow_general_knowledge=False, + json_mode=True, + context_builder_params=global_context_builder_params, + concurrent_coroutines=32, + response_type="multiple paragraphs", + ) + + logger.info("搜索引擎设置完成") + return local_search_engine, global_search_engine, local_context_builder, local_llm_params, local_context_params + + +def format_response(response): + """ + 格式化响应，添加适当的换行和段落分隔。 + """ + paragraphs = re.split(r'\n{2,}', response) + + formatted_paragraphs = [] + for para in paragraphs: + if '```' in para: + parts = para.split('```') + for i, part in enumerate(parts): + if i % 2 == 1: # 这是代码块 + parts[i] = f"\n```\n{part.strip()}\n```\n" + para = ''.join(parts) + else: + para = para.replace('. ', '.\n') + + formatted_paragraphs.append(para.strip()) + + return '\n\n'.join(formatted_paragraphs) + + +async def tavily_search(prompt: str): + """ + 使用Tavily API进行搜索 + """ + try: + client = TavilyClient(api_key=os.environ['TAVILY_API_KEY']) + resp = client.search(prompt, search_depth="advanced") + + # 将Tavily响应转换为Markdown格式 + markdown_response = "# 搜索结果\n\n" + for result in resp.get('results', []): + markdown_response += f"## [{result['title']}]({result['url']})\n\n" + markdown_response += f"{result['content']}\n\n" + + return markdown_response + except Exception as e: + raise HTTPException(status_code=500, detail=f"Tavily搜索错误: {str(e)}") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + # 启动时执行 + global local_search_engine, global_search_engine, question_generator + try: + logger.info("正在初始化搜索引擎和问题生成器...") + llm, token_encoder, text_embedder = await setup_llm_and_embedder() + entities, relationships, reports, text_units, description_embedding_store, covariates = await load_context() + local_search_engine, global_search_engine, local_context_builder, local_llm_params, local_context_params = await setup_search_engines( + llm, token_encoder, text_embedder, entities, relationships, reports, text_units, + description_embedding_store, covariates + ) + + question_generator = LocalQuestionGen( + llm=llm, + context_builder=local_context_builder, + token_encoder=token_encoder, + llm_params=local_llm_params, + context_builder_params=local_context_params, + ) + logger.info("初始化完成。") + except Exception as e: + logger.error(f"初始化过程中出错: {str(e)}") + raise + + yield + + # 关闭时执行 + logger.info("正在关闭...") + + +app = FastAPI(lifespan=lifespan) + + +# 在 chat_completions 函数中添加以下代码 + +async def full_model_search(prompt: str): + """ + 执行全模型搜索，包括本地检索、全局检索和 Tavily 搜索 + """ + local_result = await local_search_engine.asearch(prompt) + global_result = await global_search_engine.asearch(prompt) + tavily_result = await tavily_search(prompt) + + # 格式化结果 + formatted_result = "# 🔥🔥🔥综合搜索结果\n\n" + + formatted_result += "## 🔥🔥🔥本地检索结果\n" + formatted_result += format_response(local_result.response) + "\n\n" + + formatted_result += "## 🔥🔥🔥全局检索结果\n" + formatted_result += format_response(global_result.response) + "\n\n" + + formatted_result += "## 🔥🔥🔥Tavily 搜索结果\n" + formatted_result += tavily_result + "\n\n" + + return formatted_result + + +@app.post("/v1/chat/completions") +async def chat_completions(request: ChatCompletionRequest): + if not local_search_engine or not global_search_engine: + logger.error("搜索引擎未初始化") + raise HTTPException(status_code=500, detail="搜索引擎未初始化") + + try: + logger.info(f"收到聊天完成请求: {request}") + prompt = request.messages[-1].content + logger.info(f"处理提示: {prompt}") + + # 根据模型选择使用不同的搜索方法 + if request.model == "graphrag-global-search:latest": + result = await global_search_engine.asearch(prompt) + formatted_response = format_response(result.response) + elif request.model == "tavily-search:latest": + result = await tavily_search(prompt) + formatted_response = result + elif request.model == "full-model:latest": + formatted_response = await full_model_search(prompt) + else: # 默认使用本地搜索 + result = await local_search_engine.asearch(prompt) + formatted_response = format_response(result.response) + + logger.info(f"格式化的搜索结果: {formatted_response}") + + # 流式响应和非流式响应的处理保持不变 + if request.stream: + async def generate_stream(): + chunk_id = f"chatcmpl-{uuid.uuid4().hex}" + lines = formatted_response.split('\n') + for i, line in enumerate(lines): + chunk = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": request.model, + "choices": [ + { + "index": 0, + "delta": {"content": line + '\n'}, # if i > 0 else {"role": "assistant", "content": ""}, + "finish_reason": None + } + ] + } + yield f"data: {json.dumps(chunk)}\n\n" + await asyncio.sleep(0.05) + + final_chunk = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": request.model, + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop" + } + ] + } + yield f"data: {json.dumps(final_chunk)}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse(generate_stream(), media_type="text/event-stream") + else: + response = ChatCompletionResponse( + model=request.model, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=Message(role="assistant", content=formatted_response), + finish_reason="stop" + ) + ], + usage=Usage( + prompt_tokens=len(prompt.split()), + completion_tokens=len(formatted_response.split()), + total_tokens=len(prompt.split()) + len(formatted_response.split()) + ) + ) + logger.info(f"发送响应: {response}") + return JSONResponse(content=response.dict()) + + except Exception as e: + logger.error(f"处理聊天完成时出错: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + +@app.get("/v1/models") +async def list_models(): + """ + 返回可用模型列表 + """ + logger.info("收到模型列表请求") + current_time = int(time.time()) + models = [ + {"id": "graphrag-local-search:latest", "object": "model", "created": current_time - 100000, "owned_by": "graphrag"}, + {"id": "graphrag-global-search:latest", "object": "model", "created": current_time - 95000, "owned_by": "graphrag"}, + # {"id": "graphrag-question-generator:latest", "object": "model", "created": current_time - 90000, "owned_by": "graphrag"}, + # {"id": "gpt-3.5-turbo:latest", "object": "model", "created": current_time - 80000, "owned_by": "openai"}, + # {"id": "text-embedding-3-small:latest", "object": "model", "created": current_time - 70000, "owned_by": "openai"}, + {"id": "tavily-search:latest", "object": "model", "created": current_time - 85000, "owned_by": "tavily"}, + {"id": "full-model:latest", "object": "model", "created": current_time - 80000, "owned_by": "combined"} + + ] + + response = { + "object": "list", + "data": models + } + + logger.info(f"发送模型列表: {response}") + return JSONResponse(content=response) + +if __name__ == "__main__": + import uvicorn + + logger.info(f"在端口 {PORT} 上启动服务器") + uvicorn.run(app, host="0.0.0.0", port=PORT) + diff --git a/main-en.py b/main-en.py new file mode 100644 index 0000000..5f42b53 --- /dev/null +++ b/main-en.py @@ -0,0 +1,500 @@ +import os +import asyncio +import time +import uuid +import json +import re +import pandas as pd +import tiktoken +import logging +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse +from pydantic import BaseModel, Field +from typing import List, Optional, Dict, Any, Union +from contextlib import asynccontextmanager +from tavily import TavilyClient + + +# GraphRAG related imports +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.input.loaders.dfs import store_entity_semantic_embeddings +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.query.structured_search.global_search.community_context import GlobalCommunityContext +from graphrag.query.structured_search.global_search.search import GlobalSearch +from graphrag.vector_stores.lancedb import LanceDBVectorStore + +# Set up logging +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + +# Set constants and configurations +INPUT_DIR = os.getenv('INPUT_DIR') +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" +COVARIATE_TABLE = "create_final_covariates" +TEXT_UNIT_TABLE = "create_final_text_units" +COMMUNITY_LEVEL = 2 +PORT = 8012 + +# Global variables for storing search engines and question generator +local_search_engine = None +global_search_engine = None +question_generator = None + + +# Data models +class Message(BaseModel): + role: str + content: str + + +class ChatCompletionRequest(BaseModel): + model: str + messages: List[Message] + temperature: Optional[float] = 1.0 + top_p: Optional[float] = 1.0 + n: Optional[int] = 1 + stream: Optional[bool] = False + stop: Optional[Union[str, List[str]]] = None + max_tokens: Optional[int] = None + presence_penalty: Optional[float] = 0 + frequency_penalty: Optional[float] = 0 + logit_bias: Optional[Dict[str, float]] = None + user: Optional[str] = None + + +class ChatCompletionResponseChoice(BaseModel): + index: int + message: Message + finish_reason: Optional[str] = None + + +class Usage(BaseModel): + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletionResponse(BaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}") + object: str = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseChoice] + usage: Usage + system_fingerprint: Optional[str] = None + + +async def setup_llm_and_embedder(): + """ + Set up Language Model (LLM) and embedding model + """ + logger.info("Setting up LLM and embedder") + + # Get API keys and base URLs + api_key = os.environ.get("GRAPHRAG_API_KEY", "YOUR_API_KEY") + api_key_embedding = os.environ.get("GRAPHRAG_API_KEY_EMBEDDING", api_key) + api_base = os.environ.get("API_BASE", "https://api.openai.com/v1") + api_base_embedding = os.environ.get("API_BASE_EMBEDDING", "https://api.openai.com/v1") + + # Get model names + llm_model = os.environ.get("GRAPHRAG_LLM_MODEL", "gpt-3.5-turbo-0125") + embedding_model = os.environ.get("GRAPHRAG_EMBEDDING_MODEL", "text-embedding-3-small") + + # Check if API key exists + if api_key == "YOUR_API_KEY": + logger.error("Valid GRAPHRAG_API_KEY not found in environment variables") + raise ValueError("GRAPHRAG_API_KEY is not set correctly") + + # Initialize ChatOpenAI instance + llm = ChatOpenAI( + api_key=api_key, + api_base=api_base, + model=llm_model, + api_type=OpenaiApiType.OpenAI, + max_retries=20, + ) + + # Initialize token encoder + token_encoder = tiktoken.get_encoding("cl100k_base") + + # Initialize text embedding model + text_embedder = OpenAIEmbedding( + api_key=api_key_embedding, + api_base=api_base_embedding, + api_type=OpenaiApiType.OpenAI, + model=embedding_model, + deployment_name=embedding_model, + max_retries=20, + ) + + + logger.info("LLM and embedder setup complete") + return llm, token_encoder, text_embedder + + +async def load_context(): + """ + Load context data including entities, relationships, reports, text units, and covariates + """ + logger.info("Loading context data") + try: + 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="entity_description_embeddings") + description_embedding_store.connect(db_uri=LANCEDB_URI) + store_entity_semantic_embeddings(entities=entities, vectorstore=description_embedding_store) + + 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) + + covariate_df = pd.read_parquet(f"{INPUT_DIR}/{COVARIATE_TABLE}.parquet") + claims = read_indexer_covariates(covariate_df) + logger.info(f"Number of claim records: {len(claims)}") + covariates = {"claims": claims} + + logger.info("Context data loading complete") + return entities, relationships, reports, text_units, description_embedding_store, covariates + except Exception as e: + logger.error(f"Error loading context data: {str(e)}") + raise + + +async def setup_search_engines(llm, token_encoder, text_embedder, entities, relationships, reports, text_units, + description_embedding_store, covariates): + """ + Set up local and global search engines + """ + logger.info("Setting up search engines") + + # Set up local search engine + local_context_builder = LocalSearchMixedContext( + community_reports=reports, + text_units=text_units, + entities=entities, + relationships=relationships, + covariates=covariates, + 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, + "max_tokens": 12_000, + } + + local_llm_params = { + "max_tokens": 2_000, + "temperature": 0.0, + } + + local_search_engine = LocalSearch( + llm=llm, + context_builder=local_context_builder, + token_encoder=token_encoder, + llm_params=local_llm_params, + context_builder_params=local_context_params, + response_type="multiple paragraphs", + ) + + # Set up global search engine + global_context_builder = GlobalCommunityContext( + community_reports=reports, + entities=entities, + token_encoder=token_encoder, + ) + + global_context_builder_params = { + "use_community_summary": False, + "shuffle_data": True, + "include_community_rank": True, + "min_community_rank": 0, + "community_rank_name": "rank", + "include_community_weight": True, + "community_weight_name": "occurrence weight", + "normalize_community_weight": True, + "max_tokens": 12_000, + "context_name": "Reports", + } + + map_llm_params = { + "max_tokens": 1000, + "temperature": 0.0, + "response_format": {"type": "json_object"}, + } + + reduce_llm_params = { + "max_tokens": 2000, + "temperature": 0.0, + } + + global_search_engine = GlobalSearch( + llm=llm, + context_builder=global_context_builder, + token_encoder=token_encoder, + max_data_tokens=12_000, + map_llm_params=map_llm_params, + reduce_llm_params=reduce_llm_params, + allow_general_knowledge=False, + json_mode=True, + context_builder_params=global_context_builder_params, + concurrent_coroutines=32, + response_type="multiple paragraphs", + ) + + logger.info("Search engines setup complete") + return local_search_engine, global_search_engine, local_context_builder, local_llm_params, local_context_params + + +def format_response(response): + """ + Format the response by adding appropriate line breaks and paragraph separations. + """ + paragraphs = re.split(r'\n{2,}', response) + + formatted_paragraphs = [] + for para in paragraphs: + if '```' in para: + parts = para.split('```') + for i, part in enumerate(parts): + if i % 2 == 1: # This is a code block + parts[i] = f"\n```\n{part.strip()}\n```\n" + para = ''.join(parts) + else: + para = para.replace('. ', '.\n') + + formatted_paragraphs.append(para.strip()) + + return '\n\n'.join(formatted_paragraphs) + + +async def tavily_search(prompt: str): + """ + Perform a search using the Tavily API + """ + try: + client = TavilyClient(api_key=os.environ['TAVILY_API_KEY']) + resp = client.search(prompt, search_depth="advanced") + + # Convert Tavily response to Markdown format + markdown_response = "# Search Results\n\n" + for result in resp.get('results', []): + markdown_response += f"## [{result['title']}]({result['url']})\n\n" + markdown_response += f"{result['content']}\n\n" + + return markdown_response + except Exception as e: + raise HTTPException(status_code=500, detail=f"Tavily search error: {str(e)}") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + # Execute on startup + global local_search_engine, global_search_engine, question_generator + try: + logger.info("Initializing search engines and question generator...") + llm, token_encoder, text_embedder = await setup_llm_and_embedder() + entities, relationships, reports, text_units, description_embedding_store, covariates = await load_context() + local_search_engine, global_search_engine, local_context_builder, local_llm_params, local_context_params = await setup_search_engines( + llm, token_encoder, text_embedder, entities, relationships, reports, text_units, + description_embedding_store, covariates + ) + + question_generator = LocalQuestionGen( + llm=llm, + context_builder=local_context_builder, + token_encoder=token_encoder, + llm_params=local_llm_params, + context_builder_params=local_context_params, + ) + logger.info("Initialization complete.") + except Exception as e: + logger.error(f"Error during initialization: {str(e)}") + raise + + yield + + # Execute on shutdown + logger.info("Shutting down...") + + +app = FastAPI(lifespan=lifespan) + + +# Add the following code to the chat_completions function + +async def full_model_search(prompt: str): + """ + Perform a full model search, including local retrieval, global retrieval, and Tavily search + """ + local_result = await local_search_engine.asearch(prompt) + global_result = await global_search_engine.asearch(prompt) + tavily_result = await tavily_search(prompt) + + # Format results + formatted_result = "# 🔥🔥🔥Comprehensive Search Results\n\n" + + formatted_result += "## 🔥🔥🔥Local Retrieval Results\n" + formatted_result += format_response(local_result.response) + "\n\n" + + formatted_result += "## 🔥🔥🔥Global Retrieval Results\n" + formatted_result += format_response(global_result.response) + "\n\n" + + formatted_result += "## 🔥🔥🔥Tavily Search Results\n" + formatted_result += tavily_result + "\n\n" + + return formatted_result + +@app.post("/v1/chat/completions") +async def chat_completions(request: ChatCompletionRequest): + if not local_search_engine or not global_search_engine: + logger.error("Search engines not initialized") + raise HTTPException(status_code=500, detail="Search engines not initialized") + + try: + logger.info(f"Received chat completion request: {request}") + prompt = request.messages[-1].content + logger.info(f"Processing prompt: {prompt}") + + # Choose different search methods based on the model + if request.model == "graphrag-global-search:latest": + result = await global_search_engine.asearch(prompt) + formatted_response = format_response(result.response) + elif request.model == "tavily-search:latest": + result = await tavily_search(prompt) + formatted_response = result + elif request.model == "full-model:latest": + formatted_response = await full_model_search(prompt) + else: # Default to local search + result = await local_search_engine.asearch(prompt) + formatted_response = format_response(result.response) + + logger.info(f"Formatted search result: {formatted_response}") + + # Handle streaming and non-streaming responses + if request.stream: + async def generate_stream(): + chunk_id = f"chatcmpl-{uuid.uuid4().hex}" + lines = formatted_response.split('\n') + for i, line in enumerate(lines): + chunk = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": request.model, + "choices": [ + { + "index": 0, + "delta": {"content": line + '\n'}, # if i > 0 else {"role": "assistant", "content": ""}, + "finish_reason": None + } + ] + } + yield f"data: {json.dumps(chunk)}\n\n" + await asyncio.sleep(0.05) + + final_chunk = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": request.model, + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop" + } + ] + } + yield f"data: {json.dumps(final_chunk)}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse(generate_stream(), media_type="text/event-stream") + else: + response = ChatCompletionResponse( + model=request.model, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=Message(role="assistant", content=formatted_response), + finish_reason="stop" + ) + ], + usage=Usage( + prompt_tokens=len(prompt.split()), + completion_tokens=len(formatted_response.split()), + total_tokens=len(prompt.split()) + len(formatted_response.split()) + ) + ) + logger.info(f"Sending response: {response}") + return JSONResponse(content=response.dict()) + + except Exception as e: + logger.error(f"Error processing chat completion: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + +@app.get("/v1/models") +async def list_models(): + """ + Return a list of available models + """ + logger.info("Received model list request") + current_time = int(time.time()) + models = [ + {"id": "graphrag-local-search:latest", "object": "model", "created": current_time - 100000, "owned_by": "graphrag"}, + {"id": "graphrag-global-search:latest", "object": "model", "created": current_time - 95000, "owned_by": "graphrag"}, + # {"id": "graphrag-question-generator:latest", "object": "model", "created": current_time - 90000, "owned_by": "graphrag"}, + # {"id": "gpt-3.5-turbo:latest", "object": "model", "created": current_time - 80000, "owned_by": "openai"}, + # {"id": "text-embedding-3-small:latest", "object": "model", "created": current_time - 70000, "owned_by": "openai"}, + {"id": "tavily-search:latest", "object": "model", "created": current_time - 85000, "owned_by": "tavily"}, + {"id": "full-model:latest", "object": "model", "created": current_time - 80000, "owned_by": "combined"} + ] + + response = { + "object": "list", + "data": models + } + + logger.info(f"Sending model list: {response}") + return JSONResponse(content=response) + +if __name__ == "__main__": + import uvicorn + + logger.info(f"Starting server on port {PORT}") + uvicorn.run(app, host="0.0.0.0", port=PORT) diff --git a/main-graph-1.1.py b/main-graph-1.1.py new file mode 100644 index 0000000..aa81d1f --- /dev/null +++ b/main-graph-1.1.py @@ -0,0 +1,555 @@ +import os +import asyncio +import time +import uuid +import json +import re +import pandas as pd +import tiktoken +import logging +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, StreamingResponse +from pydantic import BaseModel, Field +from typing import List, Optional, Dict, Any, Union +from contextlib import asynccontextmanager +from tavily import TavilyClient +from colorama import init, Fore +# 初始化 colorama +init(autoreset=True) +# GraphRAG 相关导入 +from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey +from graphrag.query.indexer_adapters import ( + read_indexer_covariates, + read_indexer_entities, + read_indexer_communities, + read_indexer_relationships, + read_indexer_reports, + read_indexer_text_units, + read_indexer_report_embeddings +) +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.query.structured_search.global_search.community_context import GlobalCommunityContext +from graphrag.query.structured_search.global_search.search import GlobalSearch +from graphrag.query.structured_search.drift_search.drift_context import ( + DRIFTSearchContextBuilder, +) +from graphrag.config.models.drift_search_config import DRIFTSearchConfig +from graphrag.query.structured_search.drift_search.search import DRIFTSearch +from graphrag.vector_stores.lancedb import LanceDBVectorStore +# Athene-V2-Chat_exl2_2.25bpw，Rombos-Coder-V2.5-Qwen-32b-exl2_5.0bpw +LLM_MODEL = "Rombos-LLM-V2.5-Qwen-32b-4.5bpw-exl2" + +# 设置日志 +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + +# 设置常量和配置 +INPUT_DIR = "E:\\graphrag_kb\\input\\artifacts" +LANCEDB_URI = "E:\\graphrag_kb\\output\\lancedb" +COMMUNITY_REPORT_TABLE = "create_final_community_reports" +FINAL_COMMUNITY_TABLE = "create_final_communities" +ENTITY_TABLE = "create_final_nodes" +ENTITY_EMBEDDING_TABLE = "create_final_entities" +RELATIONSHIP_TABLE = "create_final_relationships" +COVARIATE_TABLE = "create_final_covariates" +TEXT_UNIT_TABLE = "create_final_text_units" +COMMUNITY_LEVEL = 10 # 社区层级，越高表示使用更精细的社区报告（但计算成本更高），默认2 +PORT = 8013 + +# 全局变量，用于存储搜索引擎和问题生成器 +local_search_engine = None +global_search_engine = None +drift_serch_engine = None +question_generator = None + + +# 数据模型 +class Message(BaseModel): + role: str + content: str + + +class ChatCompletionRequest(BaseModel): + model: str + messages: List[Message] + temperature: Optional[float] = 0.5 + top_p: Optional[float] = 0.7 + n: Optional[int] = 1 + stream: Optional[bool] = True + stop: Optional[Union[str, List[str]]] = None + max_tokens: Optional[int] = 12_000 + presence_penalty: Optional[float] = 0 + frequency_penalty: Optional[float] = 0 + logit_bias: Optional[Dict[str, float]] = None + user: Optional[str] = None + + +class ChatCompletionResponseChoice(BaseModel): + index: int + message: Message + finish_reason: Optional[str] = None + + +class Usage(BaseModel): + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletionResponse(BaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}") + object: str = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseChoice] + usage: Usage + system_fingerprint: Optional[str] = None + + +async def setup_llm_and_embedder(): + """ + 设置语言模型（LLM）和嵌入模型 + 知识图谱无法正常使用时用gemini-1.5-flash-latest模型重新训练小文档生成图谱 + """ + """ + # 获取API密钥和基础URL + api_key = "xxx" + api_key_embedding = "xxx" + api_base = "https://ai.liaobots.work/v1" + api_base_embedding = "https://ai.liaobots.work/v1" + api_base = "http://localhost:11434/v1" + # 获取模型名称 + LLM_MODEL = "gpt-4o-mini" + embedding_model = "text-embedding-ada-002" + """ + logger.info("正在设置LLM和嵌入器") + # ollama获取API密钥和基础URL + api_key = "xxx" + api_base = "http://127.0.0.1:5001/v1" + global LLM_MODEL + logger.info(Fore.CYAN + f"GRAPHRAG使用模型：{LLM_MODEL}") + # 初始化ChatOpenAI实例 + llm = ChatOpenAI( + api_key=api_key, + api_base=api_base, + model=LLM_MODEL, + api_type=OpenaiApiType.OpenAI, + max_retries=10, + request_timeout=120 # 设置超时时间为120秒 + ) + + # 初始化token编码器 + token_encoder = tiktoken.get_encoding("cl100k_base") + + # 初始化文本嵌入模型 + # openai在线模型 + """ + api_key="sk-9mxwRPHwHt8M1ct7CaCf041d6fC44e9587A041Ca3145E09e", + api_base="https://apis.wumingai.com/v1", + model=embedding_model, + deployment_name=embedding_model, + """ + # xinference本地嵌入模型 + """ + api_key="xinference", + api_base="http://127.0.0.1:9997/v1", + model="bge-m3", + deployment_name="bge-m3", + """ + text_embedder = OpenAIEmbedding( + # 本地嵌入模型 + api_key="ollama", + api_base="http://localhost:11434/v1", + model="bge-m3:Q4", + deployment_name="bge-m3:Q4", + api_type=OpenaiApiType.OpenAI, + max_retries=20, + ) + + logger.info("LLM和嵌入器设置完成") + return llm, token_encoder, text_embedder + +def embed_text(column): + text_embedder = OpenAIEmbedding( + # 本地嵌入模型 + api_key="ollama", + api_base="http://localhost:11434/v1", + model="bge-m3:Q4", + deployment_name="bge-m3:Q4", + api_type=OpenaiApiType.OpenAI, + max_retries=20, + ) + + return column.apply(lambda x: text_embedder.embed(x)) +async def load_context(): + """ + 加载上下文数据，包括实体、关系、报告、文本单元和协变量 + """ + logger.info("正在加载上下文数据") + try: + 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,content_embedding_col="full_content_embeddings") + reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL,content_embedding_col="full_content_embeddings") + full_content_embedding_store = LanceDBVectorStore(collection_name="default-community-full_content") + full_content_embedding_store.connect(db_uri=LANCEDB_URI) + read_indexer_report_embeddings(reports, full_content_embedding_store) + + final_communities_df = pd.read_parquet(f"{INPUT_DIR}/{FINAL_COMMUNITY_TABLE}.parquet") + communities = read_indexer_communities(final_communities_df,entity_df,report_df) + + + text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet") + text_units = read_indexer_text_units(text_unit_df) + covariate_df = pd.read_parquet(f"{INPUT_DIR}/{COVARIATE_TABLE}.parquet") + claims = read_indexer_covariates(covariate_df) + logger.info(f"声明记录数: {len(claims)}") + covariates = {"claims": claims} + + logger.info("上下文数据加载完成") + return entities, relationships, reports, communities,text_units, description_embedding_store, covariates + except Exception as e: + logger.error(f"加载上下文数据时出错: {str(e)}") + raise + + +async def setup_search_engines(llm, token_encoder, text_embedder, entities, relationships, reports,communities, text_units, + description_embedding_store, covariates): + """ + 设置本地搜索引擎和全局搜索引擎 + """ + logger.info("正在设置搜索引擎") + + # 设置本地搜索引擎 + local_context_builder = LocalSearchMixedContext( + community_reports=reports, + text_units=text_units, + entities=entities, + relationships=relationships, + covariates=covariates, + 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, + "max_tokens": 12_000, + } + + local_llm_params = { + "max_tokens": 12_000, + "temperature": 0.3, + } + + local_search_engine = LocalSearch( + llm=llm, + context_builder=local_context_builder, + token_encoder=token_encoder, + llm_params=local_llm_params, + context_builder_params=local_context_params, + response_type="multiple paragraphs", + ) + + # 设置全局搜索引擎 + global_context_builder = GlobalCommunityContext( + communities = communities , + community_reports=reports, + entities=entities, + token_encoder=token_encoder, + ) + + global_context_builder_params = { + "use_community_summary": False, + "shuffle_data": True, + "include_community_rank": True, + "min_community_rank": 0.5, + "community_rank_name": "rank", + "include_community_weight": True, + "community_weight_name": "occurrence weight", + "normalize_community_weight": True, + "max_tokens": 12_000, + "context_name": "Reports", + } + + map_llm_params = { + "max_tokens": 12_000, + "temperature": 0.5, + "response_format": {"type": "json_object"}, + } + + reduce_llm_params = { + "max_tokens": 12_000, + "temperature": 0.5, + } + + global_search_engine = GlobalSearch( + llm=llm, + context_builder=global_context_builder, + token_encoder=token_encoder, + max_data_tokens=12_000, + map_llm_params=map_llm_params, + reduce_llm_params=reduce_llm_params, + allow_general_knowledge=True, + json_mode=True, + context_builder_params=global_context_builder_params, + concurrent_coroutines=32, + response_type="multiple paragraphs", + ) + drift_params = DRIFTSearchConfig() + drift_params.temperature = 0.5 + drift_params.max_tokens = 12_000 + drift_context_builder = DRIFTSearchContextBuilder( + chat_llm=llm, + text_embedder=text_embedder, + entities=entities, + relationships=relationships, + reports=reports, + entity_text_embeddings=description_embedding_store, + text_units=text_units, + config = drift_params + + ) + drift_serch_engine = DRIFTSearch( + llm=llm, context_builder=drift_context_builder, token_encoder=token_encoder + ) + logger.info("搜索引擎设置完成") + return local_search_engine, global_search_engine,drift_serch_engine, local_context_builder, local_llm_params, local_context_params + + +def format_response(response): + """ + 格式化响应，添加适当的换行和段落分隔。 + """ + modified_text = re.sub(r"`(.*?)`", r'```python\1```', response) + + return modified_text + + +async def tavily_search(prompt: str): + """ + 使用Tavily API进行搜索 + """ + try: + client = TavilyClient(api_key=os.environ['TAVILY_API_KEY']) + resp = client.search(prompt, search_depth="advanced") + + # 将Tavily响应转换为Markdown格式 + markdown_response = "# 搜索结果\n\n" + for result in resp.get('results', []): + markdown_response += f"## [{result['title']}]({result['url']})\n\n" + markdown_response += f"{result['content']}\n\n" + + return markdown_response + except Exception as e: + raise HTTPException(status_code=500, detail=f"Tavily搜索错误: {str(e)}") + + +@asynccontextmanager +async def lifespan(app: FastAPI): + # 启动时执行 + global local_search_engine, global_search_engine,drift_serch_engine, question_generator + try: + logger.info("正在初始化搜索引擎和问题生成器...") + llm, token_encoder, text_embedder = await setup_llm_and_embedder() + entities, relationships, reports,communities, text_units, description_embedding_store, covariates = await load_context() + local_search_engine, global_search_engine,drift_serch_engine, local_context_builder, local_llm_params, local_context_params = await setup_search_engines( + llm, token_encoder, text_embedder, entities, relationships, reports,communities, text_units, + description_embedding_store, covariates + ) + + question_generator = LocalQuestionGen( + llm=llm, + context_builder=local_context_builder, + token_encoder=token_encoder, + llm_params=local_llm_params, + context_builder_params=local_context_params, + ) + logger.info("初始化完成。") + except Exception as e: + logger.error(f"初始化过程中出错: {str(e)}") + raise + + yield + + # 关闭时执行 + logger.info("正在关闭...") + + +app = FastAPI(lifespan=lifespan) + + +# 在 chat_completions 函数中添加以下代码 + +async def full_model_search(prompt: str): + """ + 执行全模型搜索，包括本地检索、全局检索和 Tavily 搜索 + """ + local_result = await local_search_engine.asearch(prompt) + global_result = await global_search_engine.asearch(prompt) + drift_result = await drift_serch_engine.asearch(prompt) + tavily_result = await tavily_search(prompt) + + # 格式化结果 + formatted_result = "# 🔥🔥🔥综合搜索结果\n\n" + + formatted_result += "## 🔥🔥🔥本地检索结果\n" + formatted_result += local_result.response + "\n\n" + + formatted_result += "## 🔥🔥🔥全局检索结果\n" + formatted_result += global_result.response + "\n\n" + + formatted_result += "## 🔥🔥🔥混合检索结果\n" + formatted_result += drift_result.response + "\n\n" + + formatted_result += "## 🔥🔥🔥Tavily 搜索结果\n" + formatted_result += tavily_result + "\n\n" + + return formatted_result + +@app.post("/v1/chat/completions") +async def chat_completions(request: ChatCompletionRequest): + if not local_search_engine or not global_search_engine or not drift_serch_engine: + logger.error("搜索引擎未初始化") + raise HTTPException(status_code=500, detail="搜索引擎未初始化") + + prompt = request.messages[-1].content + logger.info(Fore.CYAN + f"收到模型请求内容：{prompt}") + # 根据模型选择使用不同的搜索方法 + if request.model == "graphrag-global-search:latest": + result = await global_search_engine.asearch(prompt) + formatted_response = result.response + elif request.model == "graphrag-drift-search:latest": + result = await drift_serch_engine.asearch(prompt) + formatted_response = result.response + formatted_response:str = formatted_response["nodes"][0]["answer"] + formatted_response = formatted_response.replace(" n n","\n") + elif request.model == "tavily-search:latest": + result = await tavily_search(prompt) + formatted_response = result + elif request.model == "full-model:latest": + formatted_response = await full_model_search(prompt) + else: # 默认使用本地搜索 + result = await local_search_engine.asearch(prompt) + # 格式化回复 + #formatted_response = format_response(result.response) + formatted_response = result.response + + logger.info(Fore.CYAN + f"知识图谱的搜索结果: {formatted_response}") + # 流式响应和非流式响应的处理保持不变 + if request.stream: + async def generate_stream(): + chunk_id = f"chatcmpl-{uuid.uuid4().hex}" + lines = formatted_response.split('\n') + for i, line in enumerate(lines): + chunk = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": request.model, + "choices": [ + { + "index": 0, + #"delta": {"content": line + '\n'} if i > 0 else {"role": "assistant", "content": ""}, + "delta": {"content": line + '\n'}, # if i > 0 else {"role": "assistant", "content": ""}, + "finish_reason": None + } + ] + } + yield f"data: {json.dumps(chunk)}\n\n" + await asyncio.sleep(0.05) + + final_chunk = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": request.model, + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": "stop" + } + ] + } + yield f"data: {json.dumps(final_chunk)}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse(generate_stream(), media_type="text/event-stream") + else: + response = ChatCompletionResponse( + model=request.model, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=Message(role="assistant", content=formatted_response), + finish_reason="stop" + ) + ], + usage=Usage( + prompt_tokens=len(prompt.split()), + completion_tokens=len(formatted_response.split()), + total_tokens=len(prompt.split()) + len(formatted_response.split()) + ) + ) + logger.info(f"发送响应: {response}") + return JSONResponse(content=response.dict()) + +@app.get("/v1/models") +async def list_models(): + """ + 返回可用模型列表 + """ + logger.info("收到模型列表请求") + current_time = int(time.time()) + models = [ + {"id": "graphrag-local-search:latest", "object": "model", "created": current_time - 100000, "owned_by": "graphrag"}, + {"id": "graphrag-global-search:latest", "object": "model", "created": current_time - 95000, "owned_by": "graphrag"}, + {"id": "graphrag-drift-search:latest", "object": "model", "created": current_time - 90000, "owned_by": "graphrag"}, + # {"id": "graphrag-question-generator:latest", "object": "model", "created": current_time - 90000, "owned_by": "graphrag"}, + # {"id": "gpt-3.5-turbo:latest", "object": "model", "created": current_time - 8_0000, "owned_by": "openai"}, + # {"id": "text-embedding-3-small:latest", "object": "model", "created": current_time - 70000, "owned_by": "openai"}, + #{"id": "tavily-search:latest", "object": "model", "created": current_time - 85000, "owned_by": "tavily"}, + # {"id": "full-model:latest", "object": "model", "created": current_time - 8_0000, "owned_by": "combined"} + + ] + + response = { + "object": "list", + "data": models + } + + logger.info(f"发送模型列表: {response}") + return JSONResponse(content=response) + +if __name__ == "__main__": + import uvicorn + + logger.info(f"在端口 {PORT} 上启动服务器") + uvicorn.run(app, host="0.0.0.0", port=PORT) + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..782c9d6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,20 @@ +fastapi +uvicorn +pandas +tiktoken +graphrag +tavily-python +pydantic +python-dotenv +asyncio +aiohttp +numpy +scikit-learn +matplotlib +seaborn +nltk +spacy +transformers +torch +torchvision +torchaudio