mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-10-03 16:17:28 +00:00
refactor: RAG Refactor (#985)
Co-authored-by: Aralhi <xiaoping0501@gmail.com> Co-authored-by: csunny <cfqsunny@163.com>
This commit is contained in:
42
examples/rag/rewrite_rag_example.py
Normal file
42
examples/rag/rewrite_rag_example.py
Normal file
@@ -0,0 +1,42 @@
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from dbgpt.model import OpenAILLMClient
|
||||
from dbgpt.rag.retriever.rewrite import QueryRewrite
|
||||
|
||||
"""Query rewrite example.
|
||||
pre-requirements:
|
||||
1. install openai python sdk
|
||||
```
|
||||
pip install openai
|
||||
```
|
||||
2. set openai key and base
|
||||
```
|
||||
export OPENAI_API_KEY={your_openai_key}
|
||||
export OPENAI_API_BASE={your_openai_base}
|
||||
```
|
||||
or
|
||||
```
|
||||
import os
|
||||
os.environ["OPENAI_API_KEY"] = {your_openai_key}
|
||||
os.environ["OPENAI_API_BASE"] = {your_openai_base}
|
||||
```
|
||||
Examples:
|
||||
..code-block:: shell
|
||||
python examples/rag/rewrite_rag_example.py
|
||||
"""
|
||||
|
||||
|
||||
async def main():
|
||||
query = "compare steve curry and lebron james"
|
||||
llm_client = OpenAILLMClient()
|
||||
reinforce = QueryRewrite(
|
||||
llm_client=llm_client,
|
||||
model_name="gpt-3.5-turbo",
|
||||
)
|
||||
return await reinforce.rewrite(origin_query=query, nums=1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
output = asyncio.run(main())
|
||||
print(f"output: \n\n{output}")
|
Reference in New Issue
Block a user