chat.py

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS

import os
import gradio as gr
import json
import os

# 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)

class RepoSearch:
    def __init__(self):

        db_path = os.path.join(root_path, "database", "faiss_index")
        embeddings = OpenAIEmbeddings(api_key="sk-Mo5K9m2hKXjV1DeGeBAIzXLZFxxiOTvSwUoemKmfMXdmE9Bs", 
                              base_url="https://api.wlai.vip/v1",
                              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, 
                                          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 = ""
        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"

        return simple_str


def main():
    search = RepoSearch()

    def respond(
        prompt: str,
        history,
    ):
        if not history:
            history = [{"role": "system", "content": "You are a friendly chatbot"}]
        history.append({"role": "user", "content": prompt})

        yield history

        response = {"role": "assistant", "content": ""}
        response["content"] = search.search(prompt)
        yield history + [response]

    with gr.Blocks() as demo:
        gr.Markdown("## Semantic github search (基于语义的 github 仓库搜索) 🌐")
        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")
        prompt.submit(respond, [prompt, chatbot], [chatbot])
        prompt.submit(lambda: "", None, [prompt])

    demo.launch(share=True)


if __name__ == "__main__":
    main()