MedBridge/main.py at main · pankajarm/MedBridge · GitHub

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
"""MedBridge: Multilingual Clinical Trial Intelligence System.

Usage:
    uv run python main.py          # CLI mode
    uv run streamlit run src/ui/app.py  # Web UI mode
"""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from src.embeddings.harrier_embedder import HarrierEmbedder
from src.llm.gemma_llm import GemmaLLM
from src.storage.vector_store import VectorStore
from src.storage.graph_store import GraphStore
from src.graph.workflow import build_workflow


def main():
    print("=" * 60)
    print("  MedBridge: Multilingual Clinical Trial Intelligence")
    print("  Powered by Microsoft Harrier-OSS-v1 Embeddings")
    print("=" * 60)

    # Initialize components
    print("\nLoading models...")
    embedder = HarrierEmbedder()
    llm = GemmaLLM()

    print("Connecting to databases...")
    vector_store = VectorStore(embedding_dim=embedder.dim)
    graph_store = GraphStore()

    print(f"  Qdrant: {vector_store.count()} vectors")
    print(f"  FalkorDB: {graph_store.node_count()} nodes, {graph_store.relationship_count()} relationships")

    if vector_store.count() == 0:
        print("\nNo data found. Run setup first:")
        print("  uv run python scripts/setup_databases.py")
        return

    # Build workflow
    print("\nBuilding LangGraph workflow...")
    workflow = build_workflow(llm, embedder, vector_store, graph_store)
    print("Ready!\n")

    # Interactive CLI
    print("Enter your query (or 'quit' to exit):")
    print("Example: Find studies on metformin cardiovascular outcomes\n")

    while True:
        try:
            query = input("Query> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not query or query.lower() in ("quit", "exit", "q"):
            print("Goodbye!")
            break

        print("\nRunning multi-agent pipeline...")
        result = workflow.invoke({
            "query": query,
            "messages": [],
            "query_type": "literature_search",
            "search_results": [],
            "graph_results": [],
            "analysis_report": {},
            "agent_trace": [],
            "final_response": "",
        })

        # Print response
        print("\n" + "=" * 60)
        print(result.get("final_response", "No response generated."))
        print("=" * 60)

        # Print trace
        trace = result.get("agent_trace", [])
        if trace:
            print(f"\nAgent trace ({len(trace)} steps):")
            for i, step in enumerate(trace):
                print(f"  {i+1}. {step.get('agent', '?')} -> {step.get('action', '?')}")

        # Print search results summary
        results = result.get("search_results", [])
        if results:
            langs = {}
            for r in results:
                lang = r.get("language", "?")
                langs[lang] = langs.get(lang, 0) + 1
            print(f"\nResults: {len(results)} trials across languages: {dict(langs)}")

        print()


if __name__ == "__main__":
    main()