采用代码聊天：理解代码库的对话式人工智能 原创

1 import openai
2 import yaml
3 import os
4 import uuid
5 from qdrant_client import QdrantClient, models
6
7 # Load configuration from config.yaml
8 with open("config.yaml", "r") as file：
9 config = yaml.safe_load(file)
10
11 # Extract API keys and URLs from the config
12 qdrant_cloud_url = config["qdrant"]["url"]
13 qdrant_api_key = config["qdrant"]["api_key"]
14 openai_api_key = config["openai"]["api_key"]
15 code_folder_path = config["folder"]["path"]
16
17 # Initialize OpenAI API
18
openai.api_key = openai_api_key
19
20 # Initialize Qdrant client
21 client = QdrantClient(
22 url=qdrant_cloud_url,
23 api_key=qdrant_api_key,
24 )
25
26 def chunk_code(code, chunk_size=512)：
27 """
28 Splits the code into chunks, each of a specified size.
29 This helps in generating embeddings for manageable pieces of code.
30 """
31 lines = code.split('\n')
32 for i in range(0, len(lines), chunk_size)：
33 yield '\n'.join(lines[i：i + chunk_size])
34
35 def vectorize_and_store_code(code, filename)：
36 try：
37 # Chunk the code for better embedding representation
38 code_chunks = list(chunk_code(code))
39
40 # Generate embeddings for each chunk using the OpenAI API
41 embeddings = []
42 for chunk in code_chunks：
43 response = openai.embeddings.create(
44 input=[chunk], # Input should be a list of strings
45 model="text-embedding-ada-002"
46 )
47
48 # Access the embedding data correctly
49 embedding = response.data[0].embedding
50 embeddings.append(embedding)
51
52 # Flatten embeddings if needed or store each chunk as a separate entry
53 if len(embeddings) == 1：
54 final_embeddings = embeddings[0]
55 else：
56 final_embeddings = [item for sublist in embeddings for item in 
sublist]
57
58 # Ensure the collection exists
59 try：
60 client.create_collection(
61 collection_name="talk_to_your_code",
62 vectors_config=models.VectorParams(size=len(final_embeddings), 
distance=models.Distance.COSINE)
63 )
64 except Exception as e：
65 print("Collection already exists or other error：", e)
66
67 # Insert each chunk into the collection with relevant metadata
68 for i, embedding in enumerate(embeddings)：
69 point_id = str(uuid.uuid4())
70 points = [
71 models.PointStruct(
72 id=point_id,
73 vector=embedding,
74 payload={
75 "filename"： filename,
76 "chunk_index"： i,
77 "total_chunks"： len(embeddings),
78 "code_snippet"： code_chunks[i]
79 }
80 )
81 ]
82 client.upsert(collection_name="talk_to_your_code", points=points)
83
84 return f"{filename}： Code vectorized and stored successfully."
85
86 except Exception as e：
87 return f"An error occurred with {filename}： {str(e)}"
88
89 def process_files_in_folder(folder_path)：
90 for filename in os.listdir(folder_path)：
91 if filename.endswith(".py")：
92 file_path = os.path.join(folder_path, filename)
93 with open(file_path, 'r', encoding='utf-8') as file：
94 code = file.read()
95 print(vectorize_and_store_code(code, filename))
96
97 if __name__ == "__main__"：
98 process_files_in_folder(code_folder_path)
99

1 import chainlit as cl
2 import qdrant_client
3 import openai
4 import yaml
5 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
6 from langchain.prompts import PromptTemplate
7
8 # Load configuration from config.yaml
9 with open("config.yaml", "r") as file：
10 config = yaml.safe_load(file)
11
12 # Extract API keys and URLs from the config
13 qdrant_cloud_url = config["qdrant"]["url"]
14 qdrant_api_key = config["qdrant"]["api_key"]
15 openai_api_key = config["openai"]["api_key"]
16
17 # Initialize OpenAI API
18 openai.api_key = openai_api_key
19
20 # Initialize OpenAI Embeddings
21 embeddings = OpenAIEmbeddings(model="text-embedding-ada-002", 
openai_api_key=openai_api_key)
22
23 # Initialize Qdrant client
24 client = qdrant_client.QdrantClient(
25 url=qdrant_cloud_url,
26 api_key=qdrant_api_key,
27 )
28
29 # Initialize OpenAI Chat model
30 chat_model = ChatOpenAI(openai_api_key=openai_api_key, model="gpt-4")
31
32 # Define a simple QA prompt template
33 qa_prompt_template = PromptTemplate(
34 input_variables=["context", "question"],
35 template="Given the following context：\n{context}\nAnswer the following 
question：\n{question}"
36 )
37
38 # Chainlit function to handle user input
39 @cl.on_message
40 async def handle_message(message： cl.message.Message)：
41 try：
42 # Extract the actual text content from the message object
43 user_input = message.content
44
45 # Generate the query vector using OpenAI Embeddings
46 query_vector = embeddings.embed_query(user_input)
47
48 # Manually send the query to Qdrant
49 response = client.search(
50 collection_name="talk_to_your_code",
51 query_vector=query_vector,
52 limit=5
53 )
54
55 # Process and retrieve the relevant context (code snippets) from the 
Qdrant response
56 context_list = []
57 for point in response：
58 code_snippet = point.payload.get('code_snippet', '')
59 filename = point.payload.get('filename', 'Unknown')
60 context_list.append(f"Filename： {filename}\nCode 
Snippet：\n{code_snippet}\n")
61
62 context = "\n".join(context_list)
63 if not context：
64 context = "No matching documents found."
65
66 # Generate a response using the LLM with the retrieved context
67 prompt = qa_prompt_template.format(context=context, 
question=user_input)
68 response_text = chat_model.predict(prompt)
69
70 # Send the LLM's response
71 await cl.Message(content=response_text).send()
72
73 except Exception as e：
74 # Log the error
75 print(f"Error during message handling： {e}")
76 await cl.Message(content=f"An error occurred： {str(e)}").send()
77
78 if __name__ == "__main__"：
79 cl.run()
80