xiaoxue_admin
/
graphrag4openwebui
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

同步完整源码 - 2026-05-24

This commit is contained in:
xiaoxue 2026-05-24 23:37:54 +08:00
commit 7a596a0f63
9 changed files with 2393 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
GraphRAG4OpenWebUI/test Normal file
View File

@ -0,0 +1 @@

201
LICENSE Normal file
View File

@ -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.

161
README.md Normal file
View File

@ -0,0 +1,161 @@
### 🔥🔥🔥如有问题请联系我的微信 stoeng
### 🔥🔥🔥项目对应的视频演示请看 https://youtu.be/z4Si6O5NQ4c
# GraphRAG4OpenWebUI
<div align="center">
<p><strong>Integrate Microsoft's GraphRAG Technology into Open WebUI for Advanced Information Retrieval</strong></p>
English | <a href="README_ZH-CN.md">简体中文</a>
</div>
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)

159
README_ZH-CN.md Normal file
View File

@ -0,0 +1,159 @@
### 🔥🔥🔥如有问题请联系我的微信 stoeng
### 🔥🔥🔥项目对应的视频演示请看 https://youtu.be/z4Si6O5NQ4c
# GraphRAG4OpenWebUI
<div align="center">
<p><strong>将微软的 GraphRAG 技术集成到 Open WebUI 中，实现高级信息检索</strong></p>
<a href="README.md">English</a> | 简体中文
</div>
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)

293
graphrag3dknowledge.py Normal file
View File

@ -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}<br># 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()

503
main-cn.py Normal file
View File

@ -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)

500
main-en.py Normal file
View File

@ -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)

555
main-graph-1.1.py Normal file
View File

@ -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)

20
requirements.txt Normal file
View File

@ -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