Spaces:
Sleeping
Sleeping
| from collections.abc import Sequence | |
| from pathlib import Path | |
| import streamlit as st | |
| from neo4j.exceptions import DriverError, Neo4jError | |
| from pydantic import ValidationError | |
| from techdocs_graphrag import ( | |
| Answer, | |
| RetrievalMode, | |
| Services, | |
| SourceDocument, | |
| build_services, | |
| load_documents, | |
| ) | |
| from techdocs_graphrag.visualization import facts_to_dot | |
| SAMPLE_CORPUS = Path(__file__).parent / "data" | |
| MODE_LABELS = { | |
| RetrievalMode.GRAPH: "GraphRAG", | |
| RetrievalMode.VECTOR: "RAG vectoriel", | |
| } | |
| DISPLAY_MODES: dict[str, tuple[RetrievalMode, ...]] = { | |
| "GraphRAG": (RetrievalMode.GRAPH,), | |
| "RAG vectoriel": (RetrievalMode.VECTOR,), | |
| "Comparaison côte à côte": (RetrievalMode.VECTOR, RetrievalMode.GRAPH), | |
| } | |
| EXAMPLE_QUESTIONS = ( | |
| "Comment persister l'état d'un workflow StateGraph ?", | |
| "Quel est le rôle des reducers dans les channels ?", | |
| "Comment reprendre un workflow après une interruption humaine ?", | |
| ) | |
| st.set_page_config(page_title="TechDocs GraphRAG", page_icon="🕸️", layout="wide") | |
| def get_services() -> Services: | |
| services = build_services() | |
| services.store.verify_connectivity() | |
| services.store.ensure_schema() | |
| return services | |
| def load_services() -> Services: | |
| try: | |
| return get_services() | |
| except ValidationError as error: | |
| missing = ", ".join(sorted({str(item["loc"][0]).upper() for item in error.errors()})) | |
| st.error(f"Configuration incomplète. Variables à définir dans les secrets : {missing}") | |
| except (Neo4jError, DriverError) as error: | |
| st.error(f"Connexion à Neo4j impossible : {error}") | |
| st.stop() | |
| def index_documents(services: Services, documents: Sequence[SourceDocument]) -> None: | |
| progress = st.progress(0.0, text="Calcul des embeddings…") | |
| def report_progress(done: int, total: int) -> None: | |
| progress.progress(done / total, text=f"Extraction du graphe : {done}/{total} chunks") | |
| report = services.indexer.run(documents, on_progress=report_progress) | |
| progress.empty() | |
| st.success( | |
| f"{report.documents} documents · {report.chunks} chunks · " | |
| f"{report.entities} entités · {report.relations} relations" | |
| ) | |
| def render_sidebar(services: Services) -> tuple[RetrievalMode, ...]: | |
| with st.sidebar: | |
| st.header("Base de connaissances") | |
| uploads = st.file_uploader( | |
| "Ajouter des documents", | |
| type=["md", "markdown", "txt"], | |
| accept_multiple_files=True, | |
| ) | |
| if st.button("Indexer les documents ajoutés", disabled=not uploads): | |
| index_documents( | |
| services, | |
| [SourceDocument(name=file.name, text=file.getvalue().decode("utf-8")) for file in uploads], | |
| ) | |
| if st.button("Indexer le corpus d'exemple"): | |
| index_documents(services, load_documents(SAMPLE_CORPUS)) | |
| st.divider() | |
| display_mode = st.radio("Mode de recherche", options=list(DISPLAY_MODES)) | |
| st.divider() | |
| render_stats(services) | |
| if st.button("Effacer la conversation"): | |
| st.session_state.turns = [] | |
| st.rerun() | |
| return DISPLAY_MODES[display_mode] | |
| def render_stats(services: Services) -> None: | |
| stats = services.store.stats() | |
| first, second = st.columns(2) | |
| first.metric("Documents", stats["documents"]) | |
| second.metric("Chunks", stats["chunks"]) | |
| first.metric("Entités", stats["entities"]) | |
| second.metric("Relations", stats["relations"]) | |
| def render_answer(answer: Answer) -> None: | |
| st.markdown(answer.text) | |
| context = answer.context | |
| if context.sources: | |
| st.caption("Sources : " + " · ".join(context.sources)) | |
| with st.expander(f"Contexte · {len(context.chunks)} extraits · {len(context.facts)} relations"): | |
| graph_tab, facts_tab, chunks_tab = st.tabs(["Sous-graphe", "Relations", "Extraits"]) | |
| with graph_tab: | |
| if context.facts: | |
| st.graphviz_chart(facts_to_dot(context.facts, context.seed_entities)) | |
| else: | |
| st.info("Aucune relation du graphe n'a été utilisée.") | |
| with facts_tab: | |
| if context.facts: | |
| st.dataframe([fact.model_dump() for fact in context.facts], hide_index=True) | |
| else: | |
| st.info("Aucune relation.") | |
| with chunks_tab: | |
| for chunk in context.chunks: | |
| st.markdown(f"**{chunk.document}** · chunk {chunk.position} · score {chunk.score:.3f}") | |
| st.text(chunk.text) | |
| def render_turn(answers: Sequence[Answer]) -> None: | |
| with st.chat_message("user"): | |
| st.markdown(answers[0].question) | |
| with st.chat_message("assistant"): | |
| for column, answer in zip(st.columns(len(answers)), answers, strict=True): | |
| with column: | |
| if len(answers) > 1: | |
| st.subheader(MODE_LABELS[answer.mode]) | |
| render_answer(answer) | |
| def pick_example_question() -> str | None: | |
| st.markdown("##### Exemples de questions") | |
| columns = st.columns(len(EXAMPLE_QUESTIONS)) | |
| clicked = [ | |
| question | |
| for column, question in zip(columns, EXAMPLE_QUESTIONS, strict=True) | |
| if column.button(question) | |
| ] | |
| return clicked[0] if clicked else None | |
| def main() -> None: | |
| services = load_services() | |
| modes = render_sidebar(services) | |
| st.title("🕸️ TechDocs GraphRAG") | |
| st.caption( | |
| "Assistant de documentation technique : recherche vectorielle et graphe de connaissances " | |
| "Neo4j pour des réponses multi-hop sourcées." | |
| ) | |
| turns: list[list[Answer]] = st.session_state.setdefault("turns", []) | |
| for answers in turns: | |
| render_turn(answers) | |
| typed_question = st.chat_input("Pose une question sur la documentation…") | |
| question = typed_question or (None if turns else pick_example_question()) | |
| if not question: | |
| return | |
| with st.spinner("Recherche dans le graphe et les documents…"): | |
| answers = [services.assistant.ask(question, mode) for mode in modes] | |
| turns.append(answers) | |
| render_turn(answers) | |
| main() | |