app.py

import os
import time
import json
import asyncio
import gradio as gr

# set the env
from dotenv import load_dotenv
load_dotenv()

# get the root path of the project
current_file_path = os.path.dirname(os.path.abspath(__file__))
root_path = os.path.abspath(current_file_path)

from textwrap import dedent
from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate

class OurLLM:
    def __init__(self, model="gpt-4o"):
        '''
        params: 
            model: str, 
                模型名称 ["GLM-4-Flash", "GLM-4V-Flash", 
                         "gpt-4o-mini", "gpt-4o", "o1-mini", 
                         "gemini-1.5-flash-002", "gemini-1.5-pro-002",
                         "Qwen/Qwen2.5-7B-Instruct", "Qwen/Qwen2.5-Coder-7B-Instruct"]
        '''

        self.model_name = model

        OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
        OPENAI_API_KEY_DF = os.getenv('OPENAI_API_KEY_DF', OPENAI_API_KEY)
        OPENAI_API_KEY_AZ = os.getenv('OPENAI_API_KEY_AZ', OPENAI_API_KEY)
        OPENAI_API_KEY_CD = os.getenv('OPENAI_API_KEY_CD')
        OPENAI_API_KEY_O1 = os.getenv('OPENAI_API_KEY_O1')
        OPENAI_API_KEY_GLM = os.getenv('OPENAI_API_KEY_GLM')
        OPENAI_API_KEY_SC = os.getenv('OPENAI_API_KEY_SC')

        OPENAI_BASE_URL = os.getenv('OPENAI_BASE_URL')
        OPENAI_BASE_URL_GLM = os.getenv('OPENAI_BASE_URL_GLM')
        OPENAI_BASE_URL_SC = os.getenv('OPENAI_BASE_URL_SC')

        # 创建 API Key 映射
        apiKeyMap = {
            'gemini': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_DF},
            'gpt': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_AZ},
            'o1': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_O1},
            'claude': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_CD},
            'glm': {"base_url": OPENAI_BASE_URL_GLM, "api_key": OPENAI_API_KEY_GLM},
            'qwen': {"base_url": OPENAI_BASE_URL_SC, "api_key": OPENAI_API_KEY_SC},
        }

        for name, info in apiKeyMap.items():
            if name in model.lower():
                self.base_url = info["base_url"]
                self.api_key = info["api_key"]
                break
        assert self.base_url is not None, f"Base URL not found for model: {model}"
        assert self.api_key is not None, f"API key not found for model: {model}"

        chat_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "{system_prompt}"),
                ("human", "{input}"),
                # ("ai", "{chat_history}"),
            ]
        )
        self.chat_prompt = chat_prompt
        self.llm = self.get_llm(model)

    def clean_json(self, s):
        return s.replace("```json", "").replace("```", "").strip()

    def get_system_prompt(self, mode="assistant"):
        prompt_map = {
            "assistant": dedent("""
                你是一个智能助手，擅长用简洁的中文回答用户的问题。
                请确保你的回答准确、清晰、有条理，并且符合中文的语言习惯。
                重要提示：
                1. 回答要简洁明了，避免冗长
                2. 使用适当的专业术语
                3. 保持客观中立的语气
                4. 如果不确定，要明确指出
            """),
            # search
            "keyword_expand": dedent("""
                你是一个搜索关键词扩展专家，擅长将用户的搜索意图转化为多个相关的搜索词或短语。
                用户会输入一段描述他们搜索需求的文本，请你生成与之相关的关键词列表。
                你需要返回一个可以直接被 json 库解析的响应，包含以下内容：
                {
                    "keywords": [关键词列表],
                }
                重要提示：
                1. 关键词应该包含同义词、近义词、上位词、下位词
                2. 短语要体现不同的表达方式和组合
                3. 描述句子要涵盖不同的应用场景和用途
                4. 所有内容必须与原始搜索意图高度相关
                5. 扩展搜索意图到相关的应用场景和工具，例如:
                    - 如果搜索"PDF转MD"，应包含PDF内容提取、PDF解析工具、PDF数据处理等
                    - 如果搜索"图片压缩"，应包含批量压缩工具、图片格式转换等
                    - 如果搜索"代码格式化"，应包含代码美化工具、语法检查器、代码风格统一等
                    - 如果搜索"文本翻译"，应包含机器翻译API、多语言翻译工具、离线翻译软件等
                    - 如果搜索"数据可视化"，应包含图表生成工具、数据分析库、交互式图表等
                    - 如果搜索"网络爬虫"，应包含数据采集框架、反爬虫绕过、数据解析工具等
                    - 如果搜索"API测试"，应包含接口测试工具、性能监控、自动化测试框架等
                6. 所有内容主要使用英文表达，并对部分关键词添加额外的中文表示
                7. 返回内容不要使用任何 markdown 格式 以及任何特殊字符
            """),
            "zh2en": dedent("""
                你是一个专业的中译英翻译专家，尤其擅长学术论文的翻译工作。
                请将用户提供的中文内容翻译成地道、专业的英文。

                重要提示：
                1. 使用学术论文常用的表达方式和术语
                2. 保持专业、正式的语气
                3. 确保译文的准确性和流畅性
                4. 对专业术语进行准确翻译
                5. 遵循英文学术写作的语法规范
                6. 保持原文的逻辑结构
                7. 适当使用学术论文常见的过渡词和连接词
                8. 如遇到模糊的表达，选择最符合学术上下文的翻译
                9. 避免使用口语化或非正式的表达
                10. 注意时态和语态的准确使用
            """),
            "github_score": dedent("""
                你是一个语义匹配评分专家，擅长根据用户需求和仓库描述进行语义匹配度评分。
                用户会输入两部分内容:
                1. 用户的具体需求描述
                2. 多个仓库的描述列表(以1,2,3等数字开头)
                
                请你仔细分析用户需求，并对每个仓库进行评分。
                确保返回一个可以直接被 json 库解析的响应，包含以下内容：
                {
                    "indices": [仓库编号列表，按分数从高到低],
                    "scores": [编号对应的匹配度评分列表，0-100的整数，表示匹配程度]
                }
                
                重要提示：
                1. 评分范围为0-100的整数，高于60分表示具有明显相关性
                2. 评分要客观反映仓库与需求的契合度
                3. 只返回评分大于 60 的仓库
                4. 返回内容不要使用任何 markdown 格式 以及任何特殊字符
            """)
        }
        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. 调用 LLM 进行回答
        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):
        # 1. 获取系统提示
        if system_prompt is None:
            system_prompt = self.get_system_prompt()

        # 2. 定义异步函数
        async def call_llm(prompt):
            for _ in range(10):
                try:
                    response = await self.llm.ainvoke(prompt)
                    response.content = self.clean_json(response.content)
                    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", "init")
        embeddings = OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"), 
                                      base_url=os.getenv("OPENAI_BASE_URL"),
                                      model="text-embedding-3-small")
        
        assert os.path.exists(db_path), f"Database not found: {db_path}"
        self.vector_db = FAISS.load_local(db_path, embeddings, 
                                          index_name="init",
                                          allow_dangerous_deserialization=True)
        
    def search(self, query, k=10):
        '''
            name + description + html_url + topics
        '''
        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)
            metadata = doc.metadata
            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']}** || {content['description']}\n" # 用于大模型匹配
            simple_list.append({
                "name": content["name"],
                "description": content["description"],
                **metadata,  # 解包所有 metadata 字段
            })

        return simple_str, simple_list

def main():
    search = RepoSearch()
    llm = OurLLM(model="gpt-4o")

    def respond(
        prompt: str,
        history,
        is_llm_filter: bool = False,
        is_keyword_expand: bool = False,
        match_num: int = 40
    ):
        # 1. 初始化历史记录
        if not history:
            history = [{"role": "system", "content": "You are a friendly chatbot"}]
        history.append({"role": "user", "content": prompt})
        response = {"role": "assistant", "content": ""}
        yield history

        # 2. 扩展用户问题关键词
        if is_keyword_expand:
            response["content"] = "开始扩展关键词..."
            yield history + [response]

            query = llm.ask_question(prompt, system_prompt=llm.get_system_prompt("keyword_expand")).content
            prompt = ", ".join(json.loads(query)["keywords"])

        # 3. 语义向量匹配
        response["content"] = "开始语义向量匹配..."
        yield history + [response]
        match_str, simple_list = search.search(prompt, match_num)

        # 4. 通过 LLM 评分得到最匹配的仓库索引
        if not is_llm_filter:
            simple_strs = [f"\t**{i+1}. {repo['name']}** [✨ {repo['star_count'] // 1000}k] || **Description:** {repo['description']} || **Url:** {repo['html_url']} \n" for i, repo in enumerate(simple_list)]
            response["content"] = "".join(simple_strs)
            yield history + [response]
        else:
            response["content"] = "开始通过 LLM 评分得到最匹配的仓库..."
            yield history + [response]

            query = ' ## 用户需要的仓库内容：' + prompt + '\n ## 搜索结果列表：' + match_str
            out = llm.ask_question(query, system_prompt=llm.get_system_prompt("github_score")).content
            matched_index = json.loads(out)["indices"]

            # 5. 通过索引得到最匹配的仓库
            result = [simple_list[idx-1] for idx in matched_index]
            simple_strs = [f"\t**{i+1}. {repo['name']}** [✨ {repo['star_count'] // 1000}k] || **Description:** {repo['description']} || **Url:** {repo['html_url']} \n" for i, repo in enumerate(result)]
            response["content"] = "".join(simple_strs)
            yield history + [response]

    with gr.Blocks() as demo:
        gr.Markdown("## Github semantic search (基于语义的 github 仓库搜索) 🌐")
        
        with gr.Row():
            with gr.Column(scale=1):
                # 添加控制参数
                llm_filter = gr.Checkbox(
                    label="使用LLM过滤结果",
                    value=False,
                    info="是否使用 LLM 对搜索结果进行二次过滤"
                )
                keyword_expand = gr.Checkbox(
                    label="扩展关键词搜索",
                    value=False,
                    info="是否使用 LLM 扩展搜索关键词"
                )
                match_number = gr.Slider(
                    minimum=10,
                    maximum=100,
                    value=40,
                    step=10,
                    label="语义匹配数量",
                    info="进行语义匹配后返回的仓库数量，若使用 LLM 过滤，建议适当增加数量"
                )
            
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Agent",
                    type="messages",
                    avatar_images=(None, "https://img1.baidu.com/it/u=2193901176,1740242983&fm=253&fmt=auto&app=138&f=JPEG?w=500&h=500"),
                    height="65vh"
                )
                prompt = gr.Textbox(max_lines=2, label="Chat Message")
                
        # 更新submit调用，包含新的参数
        prompt.submit(
            respond, 
            [prompt, chatbot, llm_filter, keyword_expand, match_number], 
            [chatbot]
        )
        prompt.submit(lambda: "", None, [prompt])

    demo.launch(share=False)


if __name__ == "__main__":
    main()