codefuse-chatbot/coagent/codechat/code_search/code_search.py

# encoding: utf-8
'''
@author: 温进
@file: code_search.py
@time: 2023/11/21 下午2:35
@desc:
'''
import json
import time
from loguru import logger
from collections import defaultdict

from coagent.db_handler.graph_db_handler.nebula_handler import NebulaHandler
from coagent.db_handler.vector_db_handler.chroma_handler import ChromaHandler

from coagent.codechat.code_search.cypher_generator import CypherGenerator
from coagent.codechat.code_search.tagger import Tagger
from coagent.embeddings.get_embedding import get_embedding
from coagent.llm_models.llm_config import LLMConfig, EmbedConfig


# from configs.model_config import EMBEDDING_DEVICE, EMBEDDING_MODEL
# search_by_tag
VERTEX_SCORE = 10
HISTORY_VERTEX_SCORE = 5
VERTEX_MERGE_RATIO = 0.5

# search_by_description
MAX_DISTANCE = 1000


class CodeSearch:
    def __init__(self, llm_config: LLMConfig, nh: NebulaHandler, ch: ChromaHandler, limit: int = 3,
                 local_graph_file_path: str = ''):
        '''
        init
        @param nh: NebulaHandler
        @param ch: ChromaHandler
        @param limit: limit of result
        '''
        self.llm_config = llm_config

        self.nh = nh

        if not self.nh:
            with open(local_graph_file_path, 'r') as f:
                self.graph = json.load(f)

        self.ch = ch
        self.limit = limit

    def search_by_tag(self, query: str):
        '''
        search_code_res by tag
        @param query: str
        @return:
        '''
        tagger = Tagger()
        tag_list = tagger.generate_tag_query(query)
        logger.info(f'query tag={tag_list}')

        # get all vertices
        vertex_list = self.nh.get_vertices().get('v', [])
        vertex_vid_list = [i.as_node().get_id().as_string() for i in vertex_list]

        # update score
        vertex_score_dict = defaultdict(lambda: 0)
        for vid in vertex_vid_list:
            for tag in tag_list:
                if tag in vid:
                    vertex_score_dict[vid] += VERTEX_SCORE

        # merge depend adj score
        vertex_score_dict_final = {}
        for vertex in vertex_score_dict:
            cypher = f'''MATCH (v1)-[e]-(v2) where id(v1) == "{vertex}" RETURN v2'''
            cypher_res = self.nh.execute_cypher(cypher, self.nh.space_name)
            cypher_res_dict = self.nh.result_to_dict(cypher_res)

            adj_vertex_list = [i.as_node().get_id().as_string() for i in cypher_res_dict.get('v2', [])]

            score = vertex_score_dict.get(vertex, 0)
            for adj_vertex in adj_vertex_list:
                score += vertex_score_dict.get(adj_vertex, 0) * VERTEX_MERGE_RATIO

            if score > 0:
                vertex_score_dict_final[vertex] = score

        # get most prominent package tag
        package_score_dict = defaultdict(lambda: 0)

        for vertex, score in vertex_score_dict_final.items():
            if '#' in vertex:
                # get class name first
                cypher = f'''MATCH (v1:class)-[e:contain]->(v2) WHERE id(v2) == '{vertex}' RETURN id(v1) as id;'''
                cypher_res = self.nh.execute_cypher(cypher=cypher, format_res=True)
                class_vertices = cypher_res.get('id', [])
                if not class_vertices:
                    continue

                vertex = class_vertices[0].as_string()

            # get package name
            cypher = f'''MATCH (v1:package)-[e:contain]->(v2) WHERE id(v2) == '{vertex}' RETURN id(v1) as id;'''
            cypher_res = self.nh.execute_cypher(cypher=cypher, format_res=True)
            pac_vertices = cypher_res.get('id', [])
            if not pac_vertices:
                continue

            package = pac_vertices[0].as_string()
            package_score_dict[package] += score

        # get respective code
        res = []
        package_score_tuple = list(package_score_dict.items())
        package_score_tuple.sort(key=lambda x: x[1], reverse=True)

        ids = [i[0] for i in package_score_tuple]
        logger.info(f'ids={ids}')
        chroma_res = self.ch.get(ids=ids, include=['metadatas'])

        for vertex, score in package_score_tuple:
            index = chroma_res['result']['ids'].index(vertex)
            code_text = chroma_res['result']['metadatas'][index]['code_text']
            res.append({
                "vertex": vertex,
                "code_text": code_text}
            )
            if len(res) >= self.limit:
                break
        # logger.info(f'retrival code={res}')
        return res

    def search_by_tag_by_graph(self, query: str):
        '''
        search code by tag with graph
        @param query:
        @return:
        '''
        tagger = Tagger()
        tag_list = tagger.generate_tag_query(query)
        logger.info(f'query tag={tag_list}')

        # loop to get package node
        package_score_dict = {}
        for code, structure in self.graph.items():
            score = 0
            for class_name in structure['class_name_list']:
                for tag in tag_list:
                    if tag.lower() in class_name.lower():
                        score += 1

            for func_name_list in structure['func_name_dict'].values():
                for func_name in func_name_list:
                    for tag in tag_list:
                        if tag.lower() in func_name.lower():
                            score += 1
            package_score_dict[structure['pac_name']] = score

        # get respective code
        res = []
        package_score_tuple = list(package_score_dict.items())
        package_score_tuple.sort(key=lambda x: x[1], reverse=True)

        ids = [i[0] for i in package_score_tuple]
        logger.info(f'ids={ids}')
        chroma_res = self.ch.get(ids=ids, include=['metadatas'])

        # logger.info(chroma_res)
        for vertex, score in package_score_tuple:
            index = chroma_res['result']['ids'].index(vertex)
            code_text = chroma_res['result']['metadatas'][index]['code_text']
            res.append({
                "vertex": vertex,
                "code_text": code_text}
            )
            if len(res) >= self.limit:
                break
        # logger.info(f'retrival code={res}')
        return res

    def search_by_desciption(self, query: str, engine: str, model_path: str = "text2vec-base-chinese", embedding_device: str = "cpu", embed_config: EmbedConfig=None):
        '''
        search by perform sim search
        @param query:
        @return:
        '''
        query = query.replace(',', '，')
        query_emb = get_embedding(engine=engine, text_list=[query], model_path=model_path, embedding_device= embedding_device, embed_config=embed_config)
        query_emb = query_emb[query]

        query_embeddings = [query_emb]
        query_result = self.ch.query(query_embeddings=query_embeddings, n_results=self.limit,
                                     include=['metadatas', 'distances'])

        res = []
        for idx, distance in enumerate(query_result['result']['distances'][0]):
            if distance < MAX_DISTANCE:
                vertex = query_result['result']['ids'][0][idx]
                code_text = query_result['result']['metadatas'][0][idx]['code_text']
                res.append({
                    "vertex": vertex,
                    "code_text": code_text
                })

        return res

    def search_by_cypher(self, query: str):
        '''
        search by generating cypher
        @param query:
        @param engine:
        @return:
        '''
        cg = CypherGenerator(self.llm_config)
        cypher = cg.get_cypher(query)

        if not cypher:
            return None

        cypher_res = self.nh.execute_cypher(cypher, self.nh.space_name)
        logger.info(f'cypher execution result={cypher_res}')
        if not cypher_res.is_succeeded():
            return {
            'cypher': '',
            'cypher_res': ''
        }

        res = {
            'cypher': cypher,
            'cypher_res': cypher_res
        }

        return res


if __name__ == '__main__':
    # from configs.server_config import NEBULA_HOST, NEBULA_PORT, NEBULA_USER, NEBULA_PASSWORD, NEBULA_STORAGED_PORT
    # from configs.server_config import CHROMA_PERSISTENT_PATH
    from coagent.base_configs.env_config import (
        NEBULA_HOST, NEBULA_PORT, NEBULA_USER, NEBULA_PASSWORD, NEBULA_STORAGED_PORT,
        CHROMA_PERSISTENT_PATH
    )
    codebase_name = 'testing'

    nh = NebulaHandler(host=NEBULA_HOST, port=NEBULA_PORT, username=NEBULA_USER,
                       password=NEBULA_PASSWORD, space_name=codebase_name)
    nh.add_host(NEBULA_HOST, NEBULA_STORAGED_PORT)
    time.sleep(0.5)

    ch = ChromaHandler(path=CHROMA_PERSISTENT_PATH, collection_name=codebase_name)

    cs = CodeSearch(nh, ch)
    # res = cs.search_by_tag(tag_list=['createFineTuneCompletion', 'OpenAiApi'])
    # logger.debug(res)

    # res = cs.search_by_cypher('代码中一共有多少个类', 'openai')
    # logger.debug(res)

    res = cs.search_by_desciption('使用不同的HTTP请求类型（GET、POST、DELETE等）来执行不同的操作', 'openai')
    logger.debug(res)