370 lines
10 KiB
Python
370 lines
10 KiB
Python
from xml.parsers.expat import model
|
|
|
|
import requests
|
|
import json
|
|
import os
|
|
import time
|
|
import tqdm
|
|
import re
|
|
from dotenv import load_dotenv
|
|
import chromadb
|
|
from sentence_transformers import SentenceTransformer
|
|
import uuid
|
|
# Load variables from .env file
|
|
load_dotenv()
|
|
|
|
# --- Configuration ---
|
|
DUST_API_KEY = os.getenv("DUST_API_KEY")
|
|
WORKSPACE_ID = os.getenv("DUST_WORKSPACE_ID")
|
|
AGENT_ID = "dust" # ou l'identifiant de ton agent spécifique
|
|
|
|
BASE_URL = "https://dust.tt/api/v1"
|
|
|
|
|
|
model = SentenceTransformer("all-MiniLM-L6-v2")
|
|
HEADERS = {
|
|
"Authorization": f"Bearer {DUST_API_KEY}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
client = chromadb.PersistentClient(path="./chroma_db")
|
|
COLLECTION_NAME = "abap_rag"
|
|
collection = client.get_or_create_collection(
|
|
name=COLLECTION_NAME
|
|
)
|
|
|
|
|
|
|
|
def create_conversation(message: str) -> dict:
|
|
"""Crée une nouvelle conversation et envoie un premier message."""
|
|
url = f"{BASE_URL}/w/{WORKSPACE_ID}/assistant/conversations"
|
|
|
|
payload = {
|
|
"message": {
|
|
"content": message,
|
|
"mentions": [{"configurationId": AGENT_ID}],
|
|
"context": {
|
|
"timezone": "Europe/Paris",
|
|
"username": "python_user",
|
|
"fullName": "Python User",
|
|
"email": "python@example.com",
|
|
"profilePictureUrl": None,
|
|
},
|
|
},
|
|
"visibility": "unlisted",
|
|
"title": None,
|
|
}
|
|
|
|
response = requests.post(url, headers=HEADERS, json=payload)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
|
|
def get_agent_message_events(conversation_id: str, message_id: str):
|
|
"""Récupère la réponse de l'agent en streaming (SSE)."""
|
|
url = f"{BASE_URL}/w/{WORKSPACE_ID}/assistant/conversations/{conversation_id}/messages/{message_id}/events"
|
|
|
|
with requests.get(url, headers=HEADERS, stream=True) as response:
|
|
response.raise_for_status()
|
|
full_text = ""
|
|
|
|
for line in response.iter_lines():
|
|
if line:
|
|
decoded = line.decode("utf-8")
|
|
if decoded.startswith("data:"):
|
|
data_str = decoded[len("data:"):].strip()
|
|
try:
|
|
event = json.loads(data_str)
|
|
event_type = event.get("data").get("type")
|
|
|
|
if event_type == "generation_tokens":
|
|
token = event.get("data").get("text", "")
|
|
print(token, end="", flush=True)
|
|
full_text += token
|
|
|
|
elif event_type == "agent_message_success":
|
|
print() # Saut de ligne final
|
|
break
|
|
|
|
elif event_type == "agent_error":
|
|
print(f"\n[Erreur agent] {event}")
|
|
break
|
|
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
return full_text
|
|
|
|
|
|
def send_followup_message(conversation_id: str, message: str) -> dict:
|
|
"""Envoie un message de suivi dans une conversation existante."""
|
|
url = f"{BASE_URL}/w/{WORKSPACE_ID}/assistant/conversations/{conversation_id}/messages"
|
|
|
|
payload = {
|
|
"content": message,
|
|
"mentions": [{"configurationId": AGENT_ID}],
|
|
"context": {
|
|
"timezone": "Europe/Paris",
|
|
"username": "python_user",
|
|
"fullName": "Python User",
|
|
"email": "python@example.com",
|
|
"profilePictureUrl": None,
|
|
},
|
|
}
|
|
|
|
response = requests.post(url, headers=HEADERS, json=payload)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
|
|
def chat(message: str):
|
|
"""Point d'entrée principal : crée une conversation et affiche la réponse."""
|
|
print(f"\n[Vous] {message}")
|
|
print("[Agent] ", end="")
|
|
|
|
# 1. Créer la conversation
|
|
result = create_conversation(message)
|
|
conversation = result.get("conversation", {})
|
|
conversation_id = conversation.get("sId")
|
|
|
|
# 2. Trouver l'ID du message agent dans la conversation
|
|
agent_message = None
|
|
for msg in conversation.get("content", []):
|
|
for m in msg:
|
|
if m.get("type") == "agent_message":
|
|
agent_message = m
|
|
break
|
|
|
|
if not agent_message:
|
|
print("Aucun message agent trouvé.")
|
|
return None
|
|
|
|
agent_message_id = agent_message.get("sId")
|
|
|
|
# 3. Streamer la réponse
|
|
response_text = get_agent_message_events(conversation_id, agent_message_id)
|
|
|
|
return {"conversation_id": conversation_id, "response": response_text}
|
|
|
|
# =========================
|
|
# INDEX FILES
|
|
# =========================
|
|
|
|
def index_folder(folder_path):
|
|
|
|
all_files = []
|
|
|
|
for root, dirs, files in os.walk(folder_path):
|
|
for file in files:
|
|
|
|
if file.endswith(".abap"):
|
|
|
|
all_files.append(
|
|
os.path.join(root, file)
|
|
)
|
|
|
|
print(f"{len(all_files)} fichiers trouvés")
|
|
|
|
for filepath in tqdm.tqdm(all_files):
|
|
|
|
try:
|
|
|
|
with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
|
|
content = f.read()
|
|
|
|
chunks = split_abap_code(content)
|
|
|
|
for chunk in chunks:
|
|
|
|
embedding = model.encode(chunk).tolist()
|
|
|
|
collection.add(
|
|
ids=[str(uuid.uuid4())],
|
|
documents=[chunk],
|
|
embeddings=[embedding],
|
|
metadatas=[{
|
|
"source": filepath,
|
|
"filename": os.path.basename(filepath)
|
|
}]
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Erreur indexation {filepath}: {e}")
|
|
|
|
print("Indexation terminée")
|
|
# =========================
|
|
# DOCUMENT FILE WITH RAG
|
|
# =========================
|
|
|
|
def document_file(filepath):
|
|
|
|
with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
|
|
content = f.read()
|
|
|
|
filename = os.path.basename(filepath)
|
|
|
|
# recherche RAG
|
|
rag_context = retrieve_context(content, top_k=5)
|
|
|
|
prompt = f"""
|
|
Tu es un expert SAP ABAP.
|
|
|
|
Tu dois documenter le fichier suivant.
|
|
|
|
FICHIER ANALYSÉ:
|
|
{filename}
|
|
|
|
CHEMIN:
|
|
{filepath}
|
|
|
|
CODE PRINCIPAL:
|
|
{content}
|
|
|
|
CONTEXTE RAG:
|
|
{rag_context}
|
|
|
|
Ton résultat doit être en markdown avec plusieurs sections :
|
|
|
|
# Algorithmie
|
|
Explique le fonctionnement général du programme.
|
|
|
|
# Objets
|
|
Explique les objets utilisés et leur rôle.
|
|
|
|
# Dépendances
|
|
Explique les dépendances importantes.
|
|
|
|
# Flux de traitement
|
|
Décris les principales étapes d'exécution.
|
|
|
|
# Points d'attention
|
|
Liste les risques techniques ou métiers.
|
|
"""
|
|
|
|
return chat(prompt)
|
|
# =========================
|
|
# ABAP CHUNKER
|
|
# =========================
|
|
|
|
def split_abap_code(content):
|
|
"""
|
|
Découpe :
|
|
- METHOD
|
|
- FORM
|
|
- FUNCTION
|
|
"""
|
|
|
|
patterns = [
|
|
r"METHOD\s+.*?ENDMETHOD\.",
|
|
r"FORM\s+.*?ENDFORM\.",
|
|
r"FUNCTION\s+.*?ENDFUNCTION\."
|
|
]
|
|
|
|
chunks = []
|
|
|
|
for pattern in patterns:
|
|
matches = re.findall(
|
|
pattern,
|
|
content,
|
|
re.IGNORECASE | re.DOTALL
|
|
)
|
|
|
|
chunks.extend(matches)
|
|
|
|
# fallback si aucun chunk
|
|
if not chunks:
|
|
chunks = [content]
|
|
|
|
return chunks
|
|
# =========================
|
|
# RAG SEARCH
|
|
# =========================
|
|
|
|
def retrieve_context(query, top_k=5):
|
|
|
|
query_embedding = model.encode(query).tolist()
|
|
|
|
results = collection.query(
|
|
query_embeddings=[query_embedding],
|
|
n_results=top_k
|
|
)
|
|
|
|
documents = results["documents"][0]
|
|
metadatas = results["metadatas"][0]
|
|
|
|
context_parts = []
|
|
|
|
for doc, meta in zip(documents, metadatas):
|
|
|
|
source = meta.get("source", "unknown")
|
|
|
|
context_parts.append(f"""
|
|
SOURCE FILE:
|
|
{source}
|
|
|
|
CODE:
|
|
{doc}
|
|
""")
|
|
|
|
return "\n\n".join(context_parts)
|
|
# --- Programme principal ---
|
|
if __name__ == "__main__":
|
|
folder_or_file = input("Entrez le chemin d'un dossier ou d'un fichier à analyser : ")
|
|
outputdir = input("Entrez le chemin du dossier de sortie pour les documentations générées : ")
|
|
if os.path.isdir(folder_or_file):
|
|
all_files = []
|
|
|
|
for root, dirs, files in os.walk(folder_or_file):
|
|
for file in files:
|
|
if file.endswith(".abap"):
|
|
all_files.append(
|
|
os.path.join(root, file)
|
|
)
|
|
print("INDEXATION RAG...")
|
|
index_folder(folder_or_file)
|
|
|
|
print("GÉNÉRATION DOCUMENTATION...")
|
|
|
|
for root, dirs, files in os.walk(folder_or_file):
|
|
|
|
for file in files:
|
|
|
|
if file.endswith(".abap"):
|
|
|
|
filepath = os.path.join(root, file)
|
|
|
|
print(f"\nDocumentation de {filepath}")
|
|
|
|
result = document_file(filepath)
|
|
|
|
if result:
|
|
|
|
documentation = result['response']
|
|
|
|
output_name = os.path.splitext(file)[0]
|
|
|
|
with open(
|
|
f"{outputdir}/documentation_{output_name}.md",
|
|
"w",
|
|
encoding="utf-8"
|
|
) as f:
|
|
|
|
f.write(documentation)
|
|
|
|
print(f"Documentation générée : documentation_{output_name}.md")
|
|
elif os.path.isfile(folder_or_file):
|
|
content = open(folder_or_file, "r").read()
|
|
result = chat(f'''peux-tu m'expliquer le fonctionnement du fichier {content} ton résultat devra être en markdown avec plusieurs sections :
|
|
une section algorithmie, où tu explique le fonctionnement général du programme
|
|
une section objets, où tu explique les différents objets utilisés dans le programme et leur rôle.''' )
|
|
|
|
if result:
|
|
print(f"\n[réponse: {result}]")
|
|
print(f"\n[Conversation ID: {result['conversation_id']}]")
|
|
documentation = result['response']
|
|
file_name = os.path.basename(folder_or_file).split('.')[0]
|
|
with open(f"documentation_{file_name}.md", "w", encoding="utf-8") as f:
|
|
f.write(documentation)
|
|
else:
|
|
print("Le chemin spécifié n'est ni un fichier ni un dossier valide.")
|
|
|