Spaces:
Running
Running
File size: 12,258 Bytes
a2383a1 bb105cf 289c44a a2383a1 c084ff7 a2a0d04 a2383a1 c084ff7 a2383a1 e9d579b a2383a1 e9d579b a2383a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 |
import os
os.environ["CHAINLIT_DISABLE_WEBSOCKETS"] = "true"
# Also consider setting these for HF Spaces
os.environ["CHAINLIT_SERVER_PORT"] = "7860"
os.environ["CHAINLIT_SERVER_HOST"] = "0.0.0.0"
os.environ["CHAINLIT_USE_PREDEFINED_HOST_PORT"] = "true"
os.environ["CHAINLIT_USE_HTTP"] = "true"
import getpass
from operator import itemgetter
from typing import List, Dict
import json
import requests
#LangChain, LangGraph
from langchain_openai import ChatOpenAI
from langgraph.graph import START, StateGraph, END
from typing_extensions import List, TypedDict
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser
from langchain_core.tools import Tool, tool
from langgraph.prebuilt import ToolNode
from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
import operator
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain.vectorstores import Qdrant
from langchain.embeddings import OpenAIEmbeddings
from langchain.schema import Document
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
import chainlit as cl
import tempfile
import shutil
#helper imports
from code_analysis import *
from tools import search_pypi, write_to_docx
from prompts import describe_imports, main_prompt, documenter_prompt
from states import AgentState
if os.environ.get("SPACE_ID"): # Check if running on HF Spaces
os.environ["CHAINLIT_DISABLE_WEBSOCKETS"] = "true"
# Also consider setting these for HF Spaces
os.environ["CHAINLIT_SERVER_PORT"] = "7860"
os.environ["CHAINLIT_SERVER_HOST"] = "0.0.0.0"
# Global variables to store processed data
processed_file_path = None
document_file_path = None
vectorstore = None
main_chain = None
qdrant_client = None
@cl.on_chat_start
async def on_chat_start():
print("Chat session started")
await cl.Message(content="Welcome to the Python Code Documentation Assistant! Please upload a Python file to get started.").send()
@cl.on_message
async def on_message(message: cl.Message):
global processed_file_path, document_file_path, vectorstore, main_chain, qdrant_client
if message.elements and any(el.type == "file" for el in message.elements):
file_elements = [el for el in message.elements if el.type == "file"]
file_element = file_elements[0]
is_python_file = (
file_element.mime.startswith("text/x-python") or
file_element.name.endswith(".py") or
file_element.mime == "text/plain" # Some systems identify .py as text/plain
)
if is_python_file:
# Send processing message
msg = cl.Message(content="Processing your Python file...")
await msg.send()
print(f'file element \n {file_element} \n')
# Save uploaded file to a temporary location
temp_dir = tempfile.mkdtemp()
file_path = os.path.join(temp_dir, file_element.name)
with open(file_element.path, "rb") as source_file:
file_content_bytes = source_file.read()
with open(file_path, "wb") as destination_file:
destination_file.write(file_content_bytes)
processed_file_path = file_path
try:
# read file and extract imports
file_content = read_python_file(file_path)
imports = extract_imports(file_content, file_path)
print(f'Done reading file')
# Define describe packages graph
search_packages_tools = [search_pypi]
describe_imports_llm = ChatOpenAI(model="gpt-4o-mini")
# describe_imports_llm = describe_imports_llm.bind_tools(tools = search_packages_tools, tool_choice="required")
describe_imports_prompt = ChatPromptTemplate.from_messages([
("system", describe_imports),
("human", "{imports}")
])
describe_imports_chain = (
{"code_language": itemgetter("code_language"), "imports": itemgetter("imports")}
| describe_imports_prompt | describe_imports_llm | StrOutputParser()
)
print(f'done defining imports chain')
# Define imports chain function
def call_imports_chain(state):
last_message= state["messages"][-1]
content = json.loads(last_message.content)
chain_input = {"code_language": content['code_language'],
"imports": content['imports']}
response = describe_imports_chain.invoke(chain_input)
return {"messages": [AIMessage(content=response)]}
# bind model to tool or ToolNode
imports_tool_node = ToolNode(search_packages_tools)
# construct graph and compile
uncompiled_imports_graph = StateGraph(AgentState)
uncompiled_imports_graph.add_node("imports_agent", call_imports_chain)
uncompiled_imports_graph.add_node("imports_action", imports_tool_node)
uncompiled_imports_graph.set_entry_point("imports_agent")
def should_continue(state):
last_message = state["messages"][-1]
if last_message.tool_calls:
return "imports_action"
return END
uncompiled_imports_graph.add_conditional_edges(
"imports_agent",
should_continue
)
uncompiled_imports_graph.add_edge("imports_action", "imports_agent")
compiled_imports_graph = uncompiled_imports_graph.compile()
print(f'compiled imports graph')
# Invoke imports graph
initial_state = {
"messages": [{
"role": "human",
"content": json.dumps({
"code_language": "python",
"imports": imports
})
}]
}
# await msg.update(content="Analyzing imports and generating documentation...")
msg.content = "Analyzing your code and generating documentation..."
await msg.update()
msg = cl.Message(content="Analyzing your code and generating documentation...")
await msg.send()
result = compiled_imports_graph.invoke(initial_state)
# Define qdrant Database
qdrant_client = QdrantClient(":memory:")
embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
embedding_dim = 1536
qdrant_client.create_collection(
collection_name="description_rag_data",
vectors_config=VectorParams(size=embedding_dim, distance=Distance.COSINE),
)
vectorstore = Qdrant(qdrant_client, collection_name="description_rag_data", embeddings=embedding_model)
# Add packages chunks
text = result['messages'][-1].content
chunks = [
{"type": "Imported Packages", "name": "Imported Packages", "content": text},
#{"type": "Source Code", "name": "Source Code", "content": file_content},
]
docs = [
Document(
page_content=f"{chunk['type']} - {chunk['name']} - {chunk['content']}", # Content for the model
metadata={**chunk} # Store metadata, but don't put embeddings here
)
for chunk in chunks
]
vectorstore.add_documents(docs)
qdrant_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
print('done adding docs to DB')
#define documenter chain
documenter_llm = ChatOpenAI(model="gpt-4o-mini")
documenter_llm_prompt = ChatPromptTemplate.from_messages([
("system", documenter_prompt),
])
documenter_chain = (
{"context": itemgetter("context")}
| documenter_llm_prompt
| documenter_llm
| StrOutputParser()
)
print('done defining documenter chain')
#extract description chunks from database
collection_name = "description_rag_data"
all_points = qdrant_client.scroll(collection_name=collection_name, limit=1000)[0] # Adjust limit if needed
one_chunk = all_points[0].payload
input_text = f"type: {one_chunk['metadata']['type']} \nname: {one_chunk['metadata']['name']} \ncontent: {one_chunk['metadata']['content']}"
print('done extracting chunks form DB')
document_response = documenter_chain.invoke({"context": input_text})
print('done invoking documenter chain and will write in docx')
# write packages description in word file
document_file_path = write_to_docx(document_response)
print('done writing docx file')
# Set up Main Chain for chat
main_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
main_llm_prompt = ChatPromptTemplate.from_messages([
("system", main_prompt),
("human", "{query}")
])
main_chain = (
{"context": itemgetter("query") | qdrant_retriever, "code_language": itemgetter("code_language"), "query": itemgetter("query"), }
| main_llm_prompt
| main_llm
| StrOutputParser()
)
print('done defining main chain')
# Present download button for the document
elements = [
cl.File(
name="documentation.docx",
path=document_file_path,
display="inline"
)
]
print('done defining elements')
msg.content = "β
Your Python file has been processed! You can download the documentation file below. How can I help you with your code?"
msg.elements = elements
await msg.update()
except Exception as e:
msg.content = f"β Error processing file: {str(e)}"
await msg.update()
else:
await cl.Message(content="Please upload a Python (.py) file.").send()
# Handle chat messages if file has been processed
elif processed_file_path and main_chain:
user_input = message.content
# Send thinking message
msg = cl.Message(content="Thinking...")
await msg.send()
try:
# Use main_chain to answer the query
# invoke main chain
inputs = {
'code_language': 'Python',
'query': user_input
}
response = main_chain.invoke(inputs)
# Update with the response
msg.content = response
await msg.update()
except Exception as e:
msg.content = f"β Error processing your question: {str(e)}"
await msg.update()
else:
await cl.Message(content="Please upload a Python file first before asking questions.").send()
@cl.on_stop
def on_stop():
global processed_file_path
# Clean up temporary files
if processed_file_path and os.path.exists(os.path.dirname(processed_file_path)):
shutil.rmtree(os.path.dirname(processed_file_path))
|