add web ui

Author: billvsme

README.md

@@ -1,7 +1,7 @@
 法律AI助手
 =========
 
-法律AI助手，法律RAG，通过倒入全部200+本法律手册、网页搜索内容结合LLM回答你的问题，并且给出对应的法规和网站，基于langchain，openai，chroma，duckduckgo-search
+法律AI助手，法律RAG，通过倒入全部200+本法律手册、网页搜索内容结合LLM回答你的问题，并且给出对应的法规和网站，基于langchain，openai，chroma，duckduckgo-search, Gradio
 
 ## 原理
 
@@ -40,6 +40,14 @@ pip install -r requirements.txt
 python manager.py --init
 ```
 
+## 运行web ui
+
+```
+python manager.py --web
+```
+
+<a href="https://sm.ms/image/DbP3TiHZConUFe7" target="_blank"><img src="https://s2.loli.net/2023/10/20/DbP3TiHZConUFe7.png" ></a>
+
 ## 运行对话
 
 ```
@@ -51,5 +59,7 @@ python manager.py --shell
 ## 配置修改
 
 如果你想修改回答中的法律条数和网页条数，可以修改config.py
-- 法律条数：LAW_VS_SEARCH_K
-- 网页条数：WEB_VS_SEARCH_K
+- 法律条数: LAW_VS_SEARCH_K
+- 网页条数: WEB_VS_SEARCH_K
+- web ui地址: WEB_HOST
+- web ui端口: WEB_PORT

config.py

@@ -11,5 +11,8 @@ class Config:
     WEB_VS_COLLECTION_NAME = "web"
     WEB_VS_SEARCH_K = 2
 
+    WEB_HOST = "0.0.0.0"
+    WEB_PORT = 7860
+
 
 config = Config()

manager.py

@@ -29,7 +29,7 @@ def init_vectorstore():
     pprint(dict(info))
 
 
-async def run_shell():
+def get_law_chain():
     llm = get_llm()
     law_vs = get_vectorstore(config.LAW_VS_COLLECTION_NAME)
     web_vs = get_vectorstore(config.WEB_VS_COLLECTION_NAME)
@@ -47,6 +47,12 @@ async def run_shell():
         return_source_documents=True,
     )
 
+    return chain
+
+
+async def run_shell():
+    chain = get_law_chain()
+
     while True:
         query = input("\n用户:")
         if query.strip() == "stop":
@@ -55,8 +61,8 @@ async def run_shell():
         callback = AsyncIteratorCallbackHandler()
         task = asyncio.create_task(
             chain.ainvoke({"query": query}, config={"callbacks": [callback]}))
-        async for t in callback.aiter():
-            print(t, end="", flush=True)
+        async for new_token in callback.aiter():
+            print(new_token, end="", flush=True)
 
         print("\n")
         res = await task
@@ -65,6 +71,34 @@ async def run_shell():
         print(f"{source_text(docs)}")
 
 
+async def run_web():
+    import gradio as gr
+
+    chain = get_law_chain()
+
+    async def chat(message, history):
+        callback = AsyncIteratorCallbackHandler()
+        task = asyncio.create_task(
+            chain.ainvoke({"query": message}, config={"callbacks": [callback]}))
+
+        response = ""
+        async for new_token in callback.aiter():
+            response += new_token
+            yield response
+
+        res = await task
+        _, docs = res['result'], res['source_documents']
+
+        response += "\n" + source_text(docs)
+        yield response
+
+    demo = gr.ChatInterface(
+        fn=chat, examples=["故意杀了一个人，会判几年？", "杀人自首会减刑吗？"], title="法律AI小助手")
+
+    demo.queue()
+    demo.launch(server_name=config.WEB_HOST, server_port=config.WEB_PORT)
+
+
 if __name__ == '__main__':
     import argparse
     parser = argparse.ArgumentParser(
@@ -85,6 +119,15 @@ async def run_shell():
             run shell
         ''')
     )
+    parser.add_argument(
+        "-w",
+        "--web",
+        action="store_true",
+        help=('''
+            run web
+        ''')
+    )
+
 
     if len(sys.argv) <= 1:
         parser.print_help()
@@ -95,3 +138,5 @@ async def run_shell():
         init_vectorstore()
     if args.shell:
         asyncio.run(run_shell())
+    if args.web:
+        asyncio.run(run_web())

requirements.txt

@@ -3,3 +3,4 @@ openai==0.28.1
 chromadb==0.4.14
 tiktoken==0.5.1
 duckduckgo-search==3.9.3
+gradio==3.50.0