This is an automated email from the ASF dual-hosted git repository.
jin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-hugegraph-ai.git
The following commit(s) were added to refs/heads/main by this push:
new 759b963 feat(llm): support switch graph in api & add some query
configs (#184)
759b963 is described below
commit 759b963eabff4c421fd3352b07d15789ac9f213f
Author: PeiChaoXu <[email protected]>
AuthorDate: Mon Mar 3 16:14:48 2025 +0800
feat(llm): support switch graph in api & add some query configs (#184)
TODO: we need wrapper the query configs
---------
Co-authored-by: imbajin <[email protected]>
---
.github/workflows/hugegraph-python-client.yml | 2 +-
README.md | 2 +-
.../src/hugegraph_llm/api/models/rag_requests.py | 43 ++++++++++++++++------
hugegraph-llm/src/hugegraph_llm/api/rag_api.py | 26 ++++++++++++-
.../src/hugegraph_llm/demo/rag_demo/rag_block.py | 20 +++++++---
.../demo/rag_demo/text2gremlin_block.py | 29 ++++++++++-----
.../operators/common_op/merge_dedup_rerank.py | 12 +++---
.../src/hugegraph_llm/operators/graph_rag_task.py | 9 ++++-
.../operators/index_op/semantic_id_query.py | 6 ++-
hugegraph-ml/README.md | 28 ++++++++------
hugegraph-python-client/README.md | 4 +-
11 files changed, 131 insertions(+), 50 deletions(-)
diff --git a/.github/workflows/hugegraph-python-client.yml
b/.github/workflows/hugegraph-python-client.yml
index c0bdf5c..e05708a 100644
--- a/.github/workflows/hugegraph-python-client.yml
+++ b/.github/workflows/hugegraph-python-client.yml
@@ -20,7 +20,7 @@ jobs:
- name: Prepare HugeGraph Server Environment
run: |
docker run -d --name=graph -p 8080:8080 -e PASSWORD=admin
hugegraph/hugegraph:1.3.0
- sleep 1
+ sleep 5
- uses: actions/checkout@v4
diff --git a/README.md b/README.md
index 197161b..c0d57ae 100644
--- a/README.md
+++ b/README.md
@@ -37,7 +37,7 @@ And here are links of other repositories:
- Welcome to contribute to HugeGraph, please see
[Guidelines](https://hugegraph.apache.org/docs/contribution-guidelines/) for
more information.
- Note: It's recommended to use [GitHub Desktop](https://desktop.github.com/)
to greatly simplify the PR and commit process.
-- Code format: Please run
[`./style/code_format_and_analysis.sh`](style/code_format_and_analysis.sh) to
format your code before submitting a PR.
+- Code format: Please run
[`./style/code_format_and_analysis.sh`](style/code_format_and_analysis.sh) to
format your code before submitting a PR. (Use `pylint` to check code style)
- Thank you to all the people who already contributed to HugeGraph!
[](https://github.com/apache/incubator-hugegraph-ai/graphs/contributors)
diff --git a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
index de47aa0..9dc868a 100644
--- a/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
+++ b/hugegraph-llm/src/hugegraph_llm/api/models/rag_requests.py
@@ -23,8 +23,17 @@ from pydantic import BaseModel
from hugegraph_llm.config import prompt
+class GraphConfigRequest(BaseModel):
+ ip: str = Query('127.0.0.1', description="hugegraph client ip.")
+ port: str = Query('8080', description="hugegraph client port.")
+ name: str = Query('hugegraph', description="hugegraph client name.")
+ user: str = Query('', description="hugegraph client user.")
+ pwd: str = Query('', description="hugegraph client pwd.")
+ gs: str = None
+
+
class RAGRequest(BaseModel):
- query: str = Query("", description="Query you want to ask")
+ query: str = Query(..., description="Query you want to ask")
raw_answer: bool = Query(False, description="Use LLM to generate answer
directly")
vector_only: bool = Query(False, description="Use LLM to generate answer
with vector")
graph_only: bool = Query(True, description="Use LLM to generate answer
with graph RAG only")
@@ -33,6 +42,16 @@ class RAGRequest(BaseModel):
rerank_method: Literal["bleu", "reranker"] = Query("bleu",
description="Method to rerank the results.")
near_neighbor_first: bool = Query(False, description="Prioritize near
neighbors in the search results.")
custom_priority_info: str = Query("", description="Custom information to
prioritize certain results.")
+ # Graph Configs
+ max_graph_items: int = Query(30, description="Maximum number of items for
GQL queries in graph.")
+ topk_return_results: int = Query(20, description="Number of sorted results
to return finally.")
+ vector_dis_threshold: float = Query(0.9, description="Threshold for vector
similarity\
+ (results greater than this will be
ignored).")
+ topk_per_keyword : int = Query(1, description="TopK results returned for
each keyword \
+ extracted from the query, by default only
the most similar one is returned.")
+ client_config: Optional[GraphConfigRequest] = Query(None,
description="hugegraph server config.")
+
+ # Keep prompt params in the end
answer_prompt: Optional[str] = Query(prompt.answer_prompt,
description="Prompt to guide the answer generation.")
keywords_extract_prompt: Optional[str] = Query(
prompt.keywords_extract_prompt,
@@ -47,7 +66,18 @@ class RAGRequest(BaseModel):
# TODO: import the default value of prompt.* dynamically
class GraphRAGRequest(BaseModel):
- query: str = Query("", description="Query you want to ask")
+ query: str = Query(..., description="Query you want to ask")
+ # Graph Configs
+ max_graph_items: int = Query(30, description="Maximum number of items for
GQL queries in graph.")
+ topk_return_results: int = Query(20, description="Number of sorted results
to return finally.")
+ vector_dis_threshold: float = Query(0.9, description="Threshold for vector
similarity \
+ (results greater than this will be
ignored).")
+ topk_per_keyword : int = Query(1, description="TopK results returned for
each keyword extracted\
+ from the query, by default only the most
similar one is returned.")
+
+ client_config : Optional[GraphConfigRequest] = Query(None,
description="hugegraph server config.")
+ get_vid_only: bool = Query(False, description="return only keywords & vid
(early stop).")
+
gremlin_tmpl_num: int = Query(
1, description="Number of Gremlin templates to use. If num <=0 means
template is not provided"
)
@@ -60,15 +90,6 @@ class GraphRAGRequest(BaseModel):
)
-class GraphConfigRequest(BaseModel):
- ip: str = "127.0.0.1"
- port: str = "8080"
- name: str = "hugegraph"
- user: str = "xxx"
- pwd: str = "xxx"
- gs: str = None
-
-
class LLMConfigRequest(BaseModel):
llm_type: str
# The common parameters shared by OpenAI, Qianfan Wenxin,
diff --git a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py
b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py
index d851fd1..5e81c56 100644
--- a/hugegraph-llm/src/hugegraph_llm/api/rag_api.py
+++ b/hugegraph-llm/src/hugegraph_llm/api/rag_api.py
@@ -27,6 +27,7 @@ from hugegraph_llm.api.models.rag_requests import (
RerankerConfigRequest,
GraphRAGRequest,
)
+from hugegraph_llm.config import huge_settings
from hugegraph_llm.api.models.rag_response import RAGResponse
from hugegraph_llm.config import llm_settings, prompt
from hugegraph_llm.utils.log import log
@@ -43,6 +44,8 @@ def rag_http_api(
):
@router.post("/rag", status_code=status.HTTP_200_OK)
def rag_answer_api(req: RAGRequest):
+ set_graph_config(req)
+
result = rag_answer_func(
text=req.query,
raw_answer=req.raw_answer,
@@ -52,10 +55,15 @@ def rag_http_api(
graph_ratio=req.graph_ratio,
rerank_method=req.rerank_method,
near_neighbor_first=req.near_neighbor_first,
+ gremlin_tmpl_num=req.gremlin_tmpl_num,
+ max_graph_items=req.max_graph_items,
+ topk_return_results=req.topk_return_results,
+ vector_dis_threshold=req.vector_dis_threshold,
+ topk_per_keyword=req.topk_per_keyword,
+ # Keep prompt params in the end
custom_related_information=req.custom_priority_info,
answer_prompt=req.answer_prompt or prompt.answer_prompt,
keywords_extract_prompt=req.keywords_extract_prompt or
prompt.keywords_extract_prompt,
- gremlin_tmpl_num=req.gremlin_tmpl_num,
gremlin_prompt=req.gremlin_prompt or
prompt.gremlin_generate_prompt,
)
# TODO: we need more info in the response for users to understand the
query logic
@@ -68,16 +76,32 @@ def rag_http_api(
},
}
+ def set_graph_config(req):
+ if req.client_config:
+ huge_settings.graph_ip = req.client_config.ip
+ huge_settings.graph_port = req.client_config.port
+ huge_settings.graph_name = req.client_config.name
+ huge_settings.graph_user = req.client_config.user
+ huge_settings.graph_pwd = req.client_config.pwd
+ huge_settings.graph_space = req.client_config.gs
+
@router.post("/rag/graph", status_code=status.HTTP_200_OK)
def graph_rag_recall_api(req: GraphRAGRequest):
try:
+ set_graph_config(req)
+
result = graph_rag_recall_func(
query=req.query,
+ max_graph_items=req.max_graph_items,
+ topk_return_results=req.topk_return_results,
+ vector_dis_threshold=req.vector_dis_threshold,
+ topk_per_keyword=req.topk_per_keyword,
gremlin_tmpl_num=req.gremlin_tmpl_num,
rerank_method=req.rerank_method,
near_neighbor_first=req.near_neighbor_first,
custom_related_information=req.custom_priority_info,
gremlin_prompt=req.gremlin_prompt or
prompt.gremlin_generate_prompt,
+ get_vid_only=req.get_vid_only
)
if isinstance(result, dict):
diff --git a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py
b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py
index ba2ab7e..8261887 100644
--- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py
+++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/rag_block.py
@@ -43,6 +43,10 @@ def rag_answer(
keywords_extract_prompt: str,
gremlin_tmpl_num: Optional[int] = 2,
gremlin_prompt: Optional[str] = None,
+ max_graph_items=30,
+ topk_return_results=20,
+ vector_dis_threshold=0.9,
+ topk_per_keyword=1,
) -> Tuple:
"""
Generate an answer using the RAG (Retrieval-Augmented Generation) pipeline.
@@ -79,22 +83,28 @@ def rag_answer(
if vector_search:
rag.query_vector_index()
if graph_search:
-
rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid().import_schema(
+
rag.extract_keywords(extract_template=keywords_extract_prompt).keywords_to_vid(
+ vector_dis_threshold=vector_dis_threshold,
+ topk_per_keyword=topk_per_keyword,
+ ).import_schema(
huge_settings.graph_name
).query_graphdb(
num_gremlin_generate_example=gremlin_tmpl_num,
gremlin_prompt=gremlin_prompt,
+ max_graph_items=max_graph_items
)
# TODO: add more user-defined search strategies
rag.merge_dedup_rerank(
- graph_ratio,
- rerank_method,
- near_neighbor_first,
+ graph_ratio=graph_ratio,
+ rerank_method=rerank_method,
+ near_neighbor_first=near_neighbor_first,
+ topk_return_results=topk_return_results
)
rag.synthesize_answer(raw_answer, vector_only_answer, graph_only_answer,
graph_vector_answer, answer_prompt)
try:
- context = rag.run(verbose=True, query=text,
vector_search=vector_search, graph_search=graph_search)
+ context = rag.run(verbose=True, query=text,
vector_search=vector_search, graph_search=graph_search,
+ max_graph_items=max_graph_items)
if context.get("switch_to_bleu"):
gr.Warning("Online reranker fails, automatically switches to local
bleu rerank.")
return (
diff --git
a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py
b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py
index 46e2e9e..cce84ff 100644
--- a/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py
+++ b/hugegraph-llm/src/hugegraph_llm/demo/rag_demo/text2gremlin_block.py
@@ -188,17 +188,28 @@ def graph_rag_recall(
near_neighbor_first: bool,
custom_related_information: str,
gremlin_prompt: str,
+ max_graph_items: int,
+ topk_return_results: int,
+ vector_dis_threshold: float,
+ topk_per_keyword: int,
+ get_vid_only: bool
) -> dict:
store_schema(prompt.text2gql_graph_schema, query, gremlin_prompt)
rag = RAGPipeline()
-
-
rag.extract_keywords().keywords_to_vid().import_schema(huge_settings.graph_name).query_graphdb(
- num_gremlin_generate_example=gremlin_tmpl_num,
- gremlin_prompt=gremlin_prompt,
- ).merge_dedup_rerank(
- rerank_method=rerank_method,
- near_neighbor_first=near_neighbor_first,
- custom_related_information=custom_related_information,
- )
+ rag.extract_keywords().keywords_to_vid(
+ vector_dis_threshold=vector_dis_threshold,
+ topk_per_keyword=topk_per_keyword,
+ )
+ if not get_vid_only:
+ rag.import_schema(huge_settings.graph_name).query_graphdb(
+ num_gremlin_generate_example=gremlin_tmpl_num,
+ gremlin_prompt=gremlin_prompt,
+ max_graph_items=max_graph_items,
+ ).merge_dedup_rerank(
+ rerank_method=rerank_method,
+ near_neighbor_first=near_neighbor_first,
+ custom_related_information=custom_related_information,
+ topk_return_results=topk_return_results,
+ )
context = rag.run(verbose=True, query=query, graph_search=True)
return context
diff --git
a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py
b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py
index d9c5e98..62968fd 100644
--- a/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py
+++ b/hugegraph-llm/src/hugegraph_llm/operators/common_op/merge_dedup_rerank.py
@@ -44,7 +44,7 @@ class MergeDedupRerank:
def __init__(
self,
embedding: BaseEmbedding,
- topk: int = huge_settings.topk_return_results,
+ topk_return_results: int = huge_settings.topk_return_results,
graph_ratio: float = 0.5,
method: Literal["bleu", "reranker"] = "bleu",
near_neighbor_first: bool = False,
@@ -54,7 +54,7 @@ class MergeDedupRerank:
assert method in ["bleu", "reranker"], f"Unimplemented rerank method
'{method}'."
self.embedding = embedding
self.graph_ratio = graph_ratio
- self.topk = topk
+ self.topk_return_results = topk_return_results
self.method = method
self.near_neighbor_first = near_neighbor_first
self.custom_related_information = custom_related_information
@@ -70,11 +70,11 @@ class MergeDedupRerank:
vector_search = context.get("vector_search", False)
graph_search = context.get("graph_search", False)
if graph_search and vector_search:
- graph_length = int(self.topk * self.graph_ratio)
- vector_length = self.topk - graph_length
+ graph_length = int(self.topk_return_results * self.graph_ratio)
+ vector_length = self.topk_return_results - graph_length
else:
- graph_length = self.topk
- vector_length = self.topk
+ graph_length = self.topk_return_results
+ vector_length = self.topk_return_results
vector_result = context.get("vector_result", [])
vector_length = min(len(vector_result), vector_length)
diff --git a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py
b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py
index 7df36d7..9c3dcf9 100644
--- a/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py
+++ b/hugegraph-llm/src/hugegraph_llm/operators/graph_rag_task.py
@@ -100,12 +100,14 @@ class RAGPipeline:
by: Literal["query", "keywords"] = "keywords",
topk_per_keyword: int = huge_settings.topk_per_keyword,
topk_per_query: int = 10,
+ vector_dis_threshold: float = huge_settings.vector_dis_threshold,
):
"""
Add a semantic ID query operator to the pipeline.
:param by: Match by query or keywords.
:param topk_per_keyword: Top K results per keyword.
:param topk_per_query: Top K results per query.
+ :param vector_dis_threshold: Vector distance threshold.
:return: Self-instance for chaining.
"""
self._operators.append(
@@ -114,6 +116,7 @@ class RAGPipeline:
by=by,
topk_per_keyword=topk_per_keyword,
topk_per_query=topk_per_query,
+ vector_dis_threshold=vector_dis_threshold,
)
)
return self
@@ -174,6 +177,7 @@ class RAGPipeline:
rerank_method: Literal["bleu", "reranker"] = "bleu",
near_neighbor_first: bool = False,
custom_related_information: str = "",
+ topk_return_results: int = huge_settings.topk_return_results,
):
"""
Add a merge, deduplication, and rerank operator to the pipeline.
@@ -187,6 +191,7 @@ class RAGPipeline:
method=rerank_method,
near_neighbor_first=near_neighbor_first,
custom_related_information=custom_related_information,
+ topk_return_results=topk_return_results
)
)
return self
@@ -239,7 +244,9 @@ class RAGPipeline:
:return: Final context after all operators have been executed.
"""
if len(self._operators) == 0:
- self.extract_keywords().query_graphdb().synthesize_answer()
+ self.extract_keywords().query_graphdb(
+ max_graph_items=kwargs.get('max_graph_items')
+ ).synthesize_answer()
context = kwargs
diff --git
a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py
b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py
index 8aa6411..f1d13af 100644
--- a/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py
+++ b/hugegraph-llm/src/hugegraph_llm/operators/index_op/semantic_id_query.py
@@ -34,7 +34,8 @@ class SemanticIdQuery:
embedding: BaseEmbedding,
by: Literal["query", "keywords"] = "keywords",
topk_per_query: int = 10,
- topk_per_keyword: int = huge_settings.topk_per_keyword
+ topk_per_keyword: int = huge_settings.topk_per_keyword,
+ vector_dis_threshold: float = huge_settings.vector_dis_threshold,
):
self.index_dir = str(os.path.join(resource_path,
huge_settings.graph_name, "graph_vids"))
self.vector_index = VectorIndex.from_index_file(self.index_dir)
@@ -42,6 +43,7 @@ class SemanticIdQuery:
self.by = by
self.topk_per_query = topk_per_query
self.topk_per_keyword = topk_per_keyword
+ self.vector_dis_threshold = vector_dis_threshold
self._client = PyHugeClient(
huge_settings.graph_ip,
huge_settings.graph_port,
@@ -76,7 +78,7 @@ class SemanticIdQuery:
for keyword in keywords:
keyword_vector = self.embedding.get_text_embedding(keyword)
results = self.vector_index.search(keyword_vector,
top_k=self.topk_per_keyword,
-
dis_threshold=float(huge_settings.vector_dis_threshold))
+
dis_threshold=float(self.vector_dis_threshold))
if results:
fuzzy_match_result.extend(results[:self.topk_per_keyword])
return fuzzy_match_result
diff --git a/hugegraph-ml/README.md b/hugegraph-ml/README.md
index c6fca7f..1465db9 100644
--- a/hugegraph-ml/README.md
+++ b/hugegraph-ml/README.md
@@ -1,14 +1,16 @@
- # hugegraph-ml
+# hugegraph-ml
## Summary
-`hugegraph-ml` is a tool that integrates HugeGraph with popular graph learning
libraries.
-It implements most graph learning algorithms, enabling users to perform
end-to-end graph learning workflows directly from HugeGraph using
`hugegraph-ml`.
-Graph data can be read directly from `HugeGraph` and used for tasks such as
node embedding, node classification, and graph classification.
+`hugegraph-ml` is a tool that integrates HugeGraph with popular graph learning
libraries.
+It implements most graph learning algorithms, enabling users to perform
end-to-end graph learning workflows directly
+from HugeGraph using `hugegraph-ml`.
+Graph data can be read directly from `HugeGraph` and used for tasks such as
node embedding, node classification, and
+graph classification.
The implemented algorithm models can be found in the
[models](./src/hugegraph_ml/models) folder.
| model | paper |
-| ----------- | -------------------------------------------------- |
+|-------------|----------------------------------------------------|
| AGNN | https://arxiv.org/abs/1803.03735 |
| APPNP | https://arxiv.org/abs/1810.05997 |
| ARMA | https://arxiv.org/abs/1901.01343 |
@@ -30,13 +32,16 @@ The implemented algorithm models can be found in the
[models](./src/hugegraph_ml
## Environment Requirements
-- python 3.9+
+- python 3.9+
- hugegraph-server 1.0+
## Preparation
-1. Start the HugeGraph database, you can do it via Docker/[Binary
packages](https://hugegraph.apache.org/docs/download/download/).
- Refer to [docker-link](https://hub.docker.com/r/hugegraph/hugegraph) &
[deploy-doc](https://hugegraph.apache.org/docs/quickstart/hugegraph-server/#31-use-docker-container-convenient-for-testdev)
for guidance
+1. Start the HugeGraph database, you can do it via
+ Docker/[Binary
packages](https://hugegraph.apache.org/docs/download/download/).
+ Refer
+ to [docker-link](https://hub.docker.com/r/hugegraph/hugegraph) &
[deploy-doc](https://hugegraph.apache.org/docs/quickstart/hugegraph-server/#31-use-docker-container-convenient-for-testdev)
+ for guidance
2. Clone this project
@@ -63,7 +68,7 @@ The implemented algorithm models can be found in the
[models](./src/hugegraph_ml
### Perform node embedding on the `Cora` dataset using the `DGI` model
-Make sure that the Cora dataset is already in your HugeGraph database.
+Make sure that the Cora dataset is already in your HugeGraph database.
If not, you can run the `import_graph_from_dgl` function to import the `Cora`
dataset from `DGL` into
the `HugeGraph` database.
@@ -74,6 +79,7 @@ import_graph_from_dgl("cora")
```
Run [dgi_example.py](./src/hugegraph_ml/examples/dgi_example.py) to view the
example.
+
```bash
python ./hugegraph_ml/examples/dgi_example.py
```
@@ -112,8 +118,8 @@ embedded_graph =
node_embed_task.train_and_embed(add_self_loop=True, n_epochs=30
```python
model = MLPClassifier(
- n_in_feat=embedded_graph.ndata["feat"].shape[1],
- n_out_feat=embedded_graph.ndata["label"].unique().shape[0]
+ n_in_feat=embedded_graph.ndata["feat"].shape[1],
+ n_out_feat=embedded_graph.ndata["label"].unique().shape[0]
)
node_clf_task = NodeClassify(graph=embedded_graph, model=model)
node_clf_task.train(lr=1e-3, n_epochs=400, patience=40)
diff --git a/hugegraph-python-client/README.md
b/hugegraph-python-client/README.md
index 85383d1..31d8df7 100644
--- a/hugegraph-python-client/README.md
+++ b/hugegraph-python-client/README.md
@@ -6,10 +6,10 @@ It is used to define graph structures, perform CRUD
operations on graph data, ma
## Installation
-To install the `hugegraph-python-client`, you can use pip:
+To install the `hugegraph-python-client`, you can use pip/poetry/source
building:
```bash
-pip3 install hugegraph-python
+pip install hugegraph-python # Note: may not the latest version, recommend to
install from source
```
### Install from Source (Latest Code)
