Update app.py

app.py

@@ -10,43 +10,268 @@ import os
 current_file_path = os.path.dirname(os.path.abspath(__file__))
 root_path = os.path.abspath(current_file_path)
 
+from datetime import datetime, time
+from textwrap import dedent
+    
+from langchain_openai import ChatOpenAI
+from langchain_core.prompts import ChatPromptTemplate
+
+import re
+
+
+# 获取当前文件位置
+current_path = os.path.dirname(os.path.abspath(__file__))
+
+class OurLLM():
+    def __init__(self, model="gpt-4o"):
+
+        self.base_url = os.environ["OPENAI_BASE_URL"]
+        self.api_key = os.environ["OPENAI_API_KEY"]
+
+        # model: str, 模型名称 ["gpt-4o-mini", "gpt-4o", "o1-mini", "gemini-1.5-flash-002", "gemini-1.5-pro-002"]
+        self.big_model = "gpt-4o"
+
+        chat_prompt = ChatPromptTemplate.from_messages(
+            [
+                ("system", "{system_prompt}"),
+                ("user", "{input}"),
+            ]
+        )
+
+        self.chat_prompt = chat_prompt
+        self.llm = self.get_llm(self.big_model)
+
+    # 2. 获取指定仓库的 README 内容
+    def clean_json(self, s):
+        return s.replace("```json", "").replace("```", "")
+
+    def get_system_prompt(self, mode="assistant"):
+        prompt_map = {
+            "assistant": dedent("""
+                你是一个智能助手，擅长用简洁的中文回答用户的问题。
+                请确保你的回答准确、清晰、有条理，并且符合中文的语言习惯。
+                重要提示：
+                1. 回答要简洁明了，避免冗长
+                2. 使用适当的专业术语
+                3. 保持客观中立的语气
+                4. 如果不确定，要明确指出
+            """),
+            # paper
+            "keyword_expand": dedent("""
+                你是一个搜索关键词扩展专家，擅长将用户的搜索意图转化为多个相关的搜索词或短语。
+                用户会输入一段描述他们搜索需求的文本，请你生成与之相关的关键词列表。
+                你需要返回一个可以直接被 json 库解析的响应，不要使用任何 markdown 格式，包含以下内容：
+                {
+                    'keywords': [关键词列表],
+                }
+                重要提示：
+                1. 关键词应该包含同义词、近义词、上位词、下位词
+                2. 短语要体现不同的表达方式和组合
+                3. 描述句子要涵盖不同的应用场景和用途
+                4. 所有内容必须与原始搜索意图高度相关
+                5. 扩展搜索意图到相关的应用场景和工具，例如:
+                - 如果搜索"PDF转MD"，应包含PDF内容提取、PDF解析工具、PDF数据处理等
+                - 如果搜索"图片压缩"，应包含批量压缩工具、图片格式转换等
+                - 如果搜索"代码格式化"，应包含代码美化工具、语法检查器、代码风格统一等
+                - 如果搜索"文本翻译"，应包含机器翻译API、多语言翻译工具、离线翻译软件等
+                - 如果搜索"数据可视化"，应包含图表生成工具、数据分析库、交互式图表等
+                - 如果搜索"网络爬虫"，应包含数据采集框架、反爬虫绕过、数据解析工具等
+                - 如果搜索"API测试"，应包含接口测试工具、性能监控、自动化测试框架等
+                6. 所有内容主要使用英文表达，并对部分关键词添加额外的中文表示
+            """),
+            "github_match": dedent("""
+                你是一个仓库匹配专家，擅长根据用户需求从多个仓库中选择最合适的仓库。
+                用户会输入两部分内容:
+                1. 用户的具体需求描述
+                2. 多个仓库的描述列表(以1,2,3等数字开头)
+                
+                请你仔细分析用户需求，并从仓库列表中选择最符合需求的仓库。
+                确保返回一个可以直接被 json 库解析的响应，不要使用任何 markdown 格式，尤其是不要使用 ```json 格式，包含以下内容：
+                {
+                    "matched_repos": [匹配到的仓库编号列表，按相关度从高到低排序],
+                    "match_scores": [对应的匹配度评分列表，0-100的整数，表示匹配程度]
+                }
+                
+                重要提示：
+                1. 如果所有仓库都不相关，只返回编号1
+                2. 匹配度评分要客观反映仓库与需求的契合度
+                3. 返回的仓库数量不能超过输入仓库总数的一半
+                4. 所有内容必须使用中文表达
+            """),
+            "github_score": dedent("""
+                你是一个仓库评分专家，擅长根据用户需求对仓库进行评分。
+                用户会输入两部分内容:
+                1. 用户的具体需求描述
+                2. 多个仓库的描述列表(以1,2,3等数字开头)
+                
+                请你仔细分析用户需求，并对每个仓库进行评分。
+                确保返回一个可以直接被 json 库解析的响应，不要使用任何 markdown 格式，包含以下内容：
+                {
+                    'indices': [仓库编号列表，按分数从高到低],
+                    'scores': [编号对应的匹配度评分列表，0-100的整数，表示匹配程度]
+                }
+                
+                重要提示：
+                2. 评分范围为0-100的整数，高于60分表示具有明显相关性
+                1. 评分要客观反映仓库与需求的契合度
+                3. 只返回评分大于 60 的仓库
+            """),
+            "title_class":  dedent("""
+                你是一个内容分类专家，用户会把一些论文的题目通过多行的形式发给你，
+                请你使用尽可能简短的中文将所有题目分为不超过十类，并描述每个类别的名称以及其对应的文章数量，
+                注意你不用重复论文题目，只需要给出类别名称以及数量即可即可。
+            """),
+            "summary_struct":  dedent("""
+                你是一个论文总结专家，同时也是一个翻译专家，对中文有深入的了解，包括词汇、语法和修辞技巧，
+                能够深入分析所给英文内容的含义，可以将准备回复给用户的中文内容表示的流畅且符合中文语法习惯。
+                用户将会论文的标题以及摘要通过 json 格式发给你，请你使用尽可能简短的中文按照下列要求对所给内容进行总结：
+                回复格式为直接的 json 格式 ，样例如下：
+                {
+                    "field": "研究领域（一个词组，使用多个逗号隔开的，保证分解后的各个词语之间具有较小的重叠性）",
+                    "summary": "对论文的 abstract 进行摘要总结（保证内容尽可能的少，仅包含最关键的信息）",
+                    "translation": "将论文的 abstract 翻译为中文，保证翻译的准确性和流畅性，并忠于原文"
+                }
+                重要提示：
+                1. 直接返回JSON对象，不要添加任何其他文本、注释或标记。
+                2. 严禁在 key 对应的 value 中任何位置使用双引号！！！这样会导致我解析失败，切记！！！。
+                3. 确保你的回复可以直接通过 JSON.parse() 解析，即不要返回非 json 格式的内容和字符。
+                4. 保持回复简洁，避免重复内容。
+                5. 直接回答问题，不要重复问题内容。
+            """),
+            "field_summary": dedent("""
+                你是一个中文标签清洗专家，对中文词汇有深入了解，擅长将一对含有重复含义的标签进行合并并重命名
+                用户会将一些标签名称通过逗号隔开的形式发给你，请你将含义重复但是名称不同的标签进行合并
+                而没有含义重复的标签则不用理会，你最后需要返回你所修改的标签内容，回复为 json 格式，样例如下：
+                {
+                    "label1": "合并后的标签1",
+                    "label2": "合并后的标签1",
+                    "label3": "合并后的标签2",
+                    "label4": "合并后的标签2",
+                    "label5": "合并后的标签2"
+                }
+                重要提示：
+                1. 直接返回JSON对象, 不要添加任何其他文本、注释或标记。
+                2. 不要使用```json或任何其他格式标记。
+                3. 确保你的回复可以直接通过JSON.parse()解析。
+            """)
+        }
+        return prompt_map[mode]
+
+    def get_llm(self, model="gpt-4o-mini"):
+        '''
+        params:
+            model: str, 模型名称 ["gpt-4o-mini", "gpt-4o", "o1-mini", "gemini-1.5-flash-002"]
+        '''
+        llm = ChatOpenAI(model=model, 
+                         base_url=self.base_url, 
+                         api_key=self.api_key)
+        print(f"Init model {model} successfully!")
+        return llm
+    
+    def ask_question(self, question, system_prompt=None):
+        # 1. 获取系统提示
+        if system_prompt is None:
+            system_prompt = self.get_system_prompt()
+        
+        # 2. 生成聊天提示
+        prompt = self.chat_prompt.format(input=question, system_prompt=system_prompt)
+        config = {
+            "configurable": {"response_format": {"type": "json_object"}}
+        }
+        
+        # 3. 调用 OpenAI 模型进行回答(重调用三次，三次不成功就结束)
+        for _ in range(10):
+            try:
+                response = self.llm.invoke(prompt, config=config)
+                response.content = self.clean_json(response.content)
+                return response
+            except Exception as e:
+                print(e)
+                time.sleep(10)
+                continue
+        print(f"Failed to call llm for prompt: {prompt[0:10]}")
+        return None
+    
+    async def ask_questions_parallel(self, questions, system_prompt=None):
+        import asyncio
+        import re
+        # 1. 获取系统提示
+        if system_prompt is None:
+            system_prompt = self.get_system_prompt()
+
+        # 2. 定义异步函数
+        async def call_llm(prompt):
+            for _ in range(10):
+                try:
+                    config = {
+                        "configurable": {"response_format": {"type": "json_object"}}
+                    }
+                    response = await self.llm.ainvoke(prompt, config=config)
+                    # 1. 移除 json 标记
+                    response.content = re.sub(r'^```json\s*', '', response.content)
+                    response.content = re.sub(r'\s*```$', '', response.content)
+                    # 2. 移除公式包裹符号
+                    response.content = re.sub(r'\$', '', response.content)
+                    # 3. 移除转移符号
+                    response.content = re.sub(r'\\', '', response.content)
+                    response.content = re.sub(r'/', '', response.content)
+                    # 4. 移除各种引号
+                    response.content = re.sub(r'[“”]', '', response.content)
+                    response.content = re.sub(r'(\w+)"(\w+)', r'\1\2', response.content, flags=re.UNICODE)
+                    return response
+                except Exception as e:
+                    print(e)
+                    await asyncio.sleep(10)
+                    continue
+            print(f"Failed to call llm for prompt: {prompt[0:10]}")
+            return None
+
+        # 3. 构建 prompt
+        prompts = [self.chat_prompt.format(input=question, system_prompt=system_prompt) for question in questions]
+
+        # 4. 异步调用
+        tasks = [call_llm(prompt) for prompt in prompts]
+        results = await asyncio.gather(*tasks)
+
+        return results
+
 class RepoSearch:
     def __init__(self):
 
         db_path = os.path.join(root_path, "database", "faiss_index")
-        print(db_path)
-        # db_path = root_path
-        embeddings = OpenAIEmbeddings(api_key="sk-Mo5K9m2hKXjV1DeGeBAIzXLZFxxiOTvSwUoemKmfMXdmE9Bs", 
-                              base_url="https://api.wlai.vip/v1",
+        embeddings = OpenAIEmbeddings(api_key=os.environ["OPENAI_API_KEY"], 
+                              base_url=os.environ["OPENAI_BASE_URL"],
                               model="text-embedding-3-small")
-        print("embeddings already")
         
         assert os.path.exists(db_path), f"Database not found: {db_path}"
         self.vector_db = FAISS.load_local(db_path, embeddings, 
                                           allow_dangerous_deserialization=True)
-        print("vector_db already")
-        # pass
         
-    def search(self, query, k=20):
+    def search(self, query, k=10):
         '''
             name + description + html_url + topics
         '''
-        # return "sss"
         results = self.vector_db.similarity_search(query + " technology", k=k)
 
         simple_str = ""
+        simple_list = []
         for i, doc in enumerate(results):
             content = json.loads(doc.page_content)
             if content["description"] is None:
                 content["description"] = ""
             desc = content["description"] if len(content["description"]) < 300 else content["description"][:300] + "..."
             simple_str += f"\t**{i+1}. {content['name']}** || **Description:** {desc} || **Url:** {content['html_url']} \n"
+            simple_list.append({
+                "name": content["name"],
+                "description": desc,
+                "url": content["html_url"]
+            })
 
-        return simple_str
-
+        return simple_str, simple_list
 
 def main():
     search = RepoSearch()
+    llm = OurLLM()
 
     def respond(
         prompt: str,
@@ -59,7 +284,29 @@ def main():
         yield history
 
         response = {"role": "assistant", "content": ""}
-        response["content"] = search.search(prompt)
+        response["content"] = "开始扩展关键词..."
+        yield history + [response]
+
+        query = llm.ask_question(prompt, system_prompt=llm.get_system_prompt("keyword_expand")).content
+        json_obj = ", ".join(json.loads(query)["keywords"])
+        # response["content"] = "拓展后关键词：" + json_obj
+        # yield history + [response]
+
+        response["content"] = "开始通过 LLM 评分得到最匹配的仓库..."
+        yield history + [response]
+        simple_str, simple_list = search.search(query, 40)
+        query = json_obj + '\n' + simple_str
+        out =  llm.ask_question(query, system_prompt=llm.get_system_prompt("github_score")).content
+
+        out = out.replace('```json','').replace('```','').strip()
+        matched_repos = json.loads(out)["indices"]
+
+        result = [simple_list[idx-1] for idx in matched_repos]
+        simple_str = ""
+        for repo in result:
+            simple_str += f"\t**{repo['name']}** || **Description:** {repo['description']} || **Url:** {repo['url']} \n"
+        response = {"role": "assistant", "content": ""}
+        response["content"] = simple_str
         yield history + [response]
 
     with gr.Blocks() as demo: