Rename backend.py to app.py

app.py

@@ -0,0 +1,596 @@
+import sqlite3
+import contextlib
+import json
+from http.server import BaseHTTPRequestHandler
+from urllib.parse import urlparse, parse_qs
+import traceback
+from pydantic import BaseModel, Field
+from typing import List, Dict, Tuple
+import os
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import FakeEmbeddings
+from langchain_community.vectorstores.utils import DistanceStrategy
+from together import Together
+import numpy as np
+from collections import defaultdict
+
+app = FastAPI(title="Knowledge Graph API")
+
+# Enable CORS for frontend access
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Database configuration - UPDATE THESE PATHS
+DATABASE_CONFIG = {
+    "triplets_db": "triplets_new.db",  
+    "definitions_db": "relations_new.db",
+    "news_db": "cnnhealthnews2.db",
+    "triplets_table": "triplets", 
+    "definitions_table": "relations", 
+    "head_column": "head_entity",
+    "relation_column": "relation", 
+    "tail_column": "tail_entity",
+    "definition_column": "definition",
+    "link_column": "link",
+    "title_column": "column",
+    "content_column": "content"
+}
+
+class GraphNode(BaseModel):
+    id: str
+    label: str
+    type: str = "entity"
+
+class GraphEdge(BaseModel):
+    source: str
+    target: str
+    relation: str
+    definition: Optional[str] = None
+
+class GraphData(BaseModel):
+    nodes: List[GraphNode]
+    edges: List[GraphEdge]
+
+class TripletData(BaseModel):
+    head: str
+    relation: str
+    tail: str
+
+class RelationDefinition(BaseModel):
+    relation: str
+    definition: str
+
+class RetrieveTripletsResponse(BaseModel):
+    triplets: List[TripletData]
+    relations: List[RelationDefinition]
+
+class NewsItem(BaseModel):
+    url: str
+    content: str
+    preview: str
+    title: str
+
+class QueryRequest(BaseModel):
+    query: str
+    
+class QueryResponse(BaseModel):
+    answer: str
+    triplets: List[TripletData]
+    relations: List[RelationDefinition]
+    news_items: List[NewsItem]
+    graph_data: GraphData
+    
+class ExtractedInformationNews(BaseModel):
+    extracted_information: str = Field(description="Extracted information")
+    links: list = Field(description="citation links")
+
+class ExtractedInformation(BaseModel):
+    extracted_information: str = Field(description="Extracted information")
+
+@contextlib.contextmanager
+def get_triplets_db():
+    conn = None
+    try:
+        conn = sqlite3.connect(DATABASE_CONFIG["triplets_db"])
+        yield conn
+    finally:
+        if conn:
+            conn.close()
+
+@contextlib.contextmanager
+def get_news_db():
+    conn = None
+    try:
+        conn = sqlite3.connect(DATABASE_CONFIG["news_db"])
+        yield conn
+    finally:
+        if conn:
+            conn.close()
+
+@contextlib.contextmanager
+def get_definitions_db():
+    conn = None
+    try:
+        conn = safe_connect(DATABASE_CONFIG["definitions_db"])
+        yield conn
+    finally:
+        if conn:
+            conn.close()
+
+def retrieve_triplets(query: str) -> Tuple[List[Tuple[str, str, str]], List[Tuple[str, str]]]:
+    """
+    Args:
+        query (str): User query
+        
+    Returns:
+        Tuple containing:
+        - List of triplets: [(head, relation, tail), ...]
+        - List of relations with definitions: [(relation, definition), ...]
+    """
+    API_KEY = os.environ.get("TOGETHER_API_KEY")
+    client = Together(api_key = API_KEY)
+    
+    dummy_embeddings = FakeEmbeddings(size=768)
+    triplets_store = FAISS.load_local(
+    "triplets_index_compressed", dummy_embeddings, allow_dangerous_deserialization=True
+    )
+    triplets_store.index.nprobe = 100
+    triplets_store._normalize_L2 = True
+    triplets_store.distance_strategy = DistanceStrategy.COSINE
+
+    response = client.embeddings.create(
+      model = "Alibaba-NLP/gte-modernbert-base",
+      input = query
+    )
+
+    emb = np.array(response.data[0].embedding)
+    emb = emb / np.linalg.norm(emb)
+    
+    related_head_entity = []
+    result_triplets = triplets_store.similarity_search_with_score_by_vector(emb, k=100)
+    for res, score in result_triplets:
+        if score > 0.7:
+            related_head_entity.append(res)
+            
+    try:
+        all_triplets = []
+        with get_triplets_db() as conn:
+            head_col = DATABASE_CONFIG["head_column"]
+            rel_col = DATABASE_CONFIG["relation_column"] 
+            tail_col = DATABASE_CONFIG["tail_column"]
+        
+            for head_entity in related_head_entity:
+                he = head_entity.page_content
+                cursor = conn.cursor()
+                cursor.execute("SELECT * FROM triplets WHERE head_entity = (?)", ([he]))
+                rows = cursor.fetchall()
+                triplets = [(str(row[0]), str(row[1]), str(row[2])) for row in rows]
+                all_triplets += triplets
+            
+        all_relations = []
+        relations = [relation for _, relation, _ in all_triplets]
+        with get_definitions_db() as conn:
+            rel_col = DATABASE_CONFIG["relation_column"] 
+            def_col = DATABASE_CONFIG["definition_column"]
+        
+            for rel in set(relations):
+                cursor = conn.cursor()
+                cursor.execute("SELECT * FROM relations WHERE relation = (?)", ([rel]))
+                rows = cursor.fetchall()
+                relation = [(str(row[0]), str(row[1])) for row in rows]
+                all_relations += relation
+
+        return all_triplets, all_relations
+        
+    except Exception as e:
+        print(f"Error in retrieve_triplets: {e}")
+        return [], []
+
+def retrieve_news(query: str) -> Dict[str, str]:
+    """
+    Args:
+        query (str): User query
+        
+    Returns: Tuple
+        - Related content
+        - Links of the related content
+    """
+    API_KEY = os.environ.get("TOGETHER_API_KEY")
+    client = Together(api_key = API_KEY)
+    
+    dummy_embeddings = FakeEmbeddings(size=768)
+    news_store = FAISS.load_local(
+        "news_index_compressed", dummy_embeddings, allow_dangerous_deserialization=True
+    )
+    news_store.index.nprobe = 100
+    news_store._normalize_L2 = True
+    news_store.distance_strategy = DistanceStrategy.COSINE
+    
+    news_store._normalize_L2 = True
+    news_store.distance_strategy = DistanceStrategy.COSINE
+
+    response = client.embeddings.create(
+      model = "Alibaba-NLP/gte-modernbert-base",
+      input = query
+    )
+
+    emb = np.array(response.data[0].embedding)
+    emb = emb / np.linalg.norm(emb)
+
+    related_news_content = []
+    result_news= news_store.similarity_search_with_score_by_vector(emb, k=500)
+    for res, score in result_news:
+        if score > 0.7:
+            print(score)
+            related_news_content.append(res)
+    
+    news_dict = defaultdict(list)
+    links = [res.metadata["link"] for res in related_news_content]
+    for idx, link in enumerate(links):
+        news_dict[link].append(related_news_content[idx].page_content)
+        
+    content_only = [". ".join(sentences) for sentences in news_dict.values()]
+    
+    return content_only, links
+
+
+def extract_information_from_triplets(query: str,
+                                      triplets: List[Tuple[str, str, str]], 
+                                      relations: List[Tuple[str, str]]) -> str:
+    """
+    REPLACE THIS FUNCTION WITH YOUR ACTUAL IMPLEMENTATION
+    
+    Args:
+        triplets: List of triplets from retrieve_triplets
+        relations: List of relation definitions from retrieve_triplets
+        
+    Returns:
+        str: Extracted information from triplets
+    """
+    system_prompt = f'''Given a a list of relational triplets and a list of relation and its definition. Extract the information from the triplets to answer query question.
+    If there is no related or useful information can be extracted from the triplets to answer the query question, inform "No related information found."
+    Give the output in paragraphs form narratively, you can explain the reason behind your answer in detail."
+    '''
+    
+    user_prompt = f'''
+    query question: {query}
+    list of triplets: {triplets}
+    list of relations and their definition: {relations}
+    extracted information:
+    '''
+
+    API_KEY = os.environ.get("TOGETHER_API_KEY")
+    client = Together(api_key = API_KEY)
+
+    response = client.chat.completions.create(
+        model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
+        temperature = 0,
+        messages=[{
+            "role": "system",
+            "content": [
+                {"type": "text", "text":system_prompt}
+            ]
+        },
+                {
+           "role": "user",
+           "content": [
+               {"type": "text", "text":user_prompt},
+            ]
+        }]
+    )
+    
+    return response.choices[0].message.content
+
+def extract_information_from_news(query: str,
+                                  news_list: Dict[str, str]) -> Tuple[str, List[str]]:
+   """   
+    Args:
+        news_list: List from retrieve_news
+        
+    Returns:
+        Extracted information string
+    """
+    system_prompt = f'''Given a list of some information related to the query, extract all important information from the list to answer query question.
+    Every item in the list represent one information, if the information is ambiguous (e.g. contains unknown pronoun to which it refers), do not use that information to answer the query.
+    You don't have to use all the information, only use the information that has clarity and a good basis, but try to use as many information as possible.
+    If there is no related or useful information can be extracted from the news information to answer the query question, write "No related information found." as the extracted_information output.
+    Give the extracted_information output in paragraphs form detailedly.
+    The output must be in this form: {{"extracted_information": <output paragraphs>}}
+    '''
+    
+    user_prompt = f'''
+    query: {query}
+    news list: {news_list}
+    output:
+    '''
+
+    response = client.chat.completions.create(
+       model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
+       response_format={
+            "type": "json_schema",
+            "schema": ExtractedInformation.model_json_schema(),
+       },
+       temperature = 0,
+       messages=[{
+           "role": "system",
+           "content": [
+               {"type": "text", "text":system_prompt}
+           ]
+       },
+                {
+           "role": "user",
+           "content": [
+               {"type": "text", "text":user_prompt},
+           ]
+       }]
+    )
+    response = json.loads(response.choices[0].message.content)
+    info = response['extracted_information']
+    
+    return info
+
+def extract_information(query:str, triplet_info: str, news_info: str, language:str) -> str:
+    """   
+    Args:
+        triplet_info: Information extracted from triplets
+        news_info: Information extracted from news
+        
+    Returns:
+        str: Final answer for the user
+    """
+    client = Together(api_key = API_KEY)
+    system_prompt = f'''Given information from two sources, combine the information and make a comprehensive and informative paragraph that answer the query.
+    Make sure the output paragraph includes all crucial information and given in detail.
+    If there is no related or useful information can be extracted from the triplets to answer the query question, inform "No related information found."
+    Remember this paragraph will be shown to user, so make sure it is based on facts and data, also use appropriate language.
+    The output must be in this form and in {language} language: {{"extracted_information": <output paragraphs>}}
+    '''
+    
+    user_prompt = f'''
+    query: {query}
+    first source: {triplet_info}
+    second source: {news_info}
+    extracted information:
+    '''
+
+    response = client.chat.completions.create(
+       model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
+       response_format={
+            "type": "json_schema",
+            "schema": ExtractedInformation.model_json_schema(),
+       },
+       temperature = 0,
+       messages=[{
+           "role": "system",
+           "content": [
+               {"type": "text", "text":system_prompt}
+           ]
+       },
+                {
+           "role": "user",
+           "content": [
+               {"type": "text", "text":user_prompt},
+           ]
+       }]
+    )
+    
+    response = json.loads(response.choices[0].message.content)
+    answer = response["extracted_information"]
+    return answer
+
+def news_preview(links: list[str]) -> Tuple[str, str, str]:
+    try:
+        preview_contents = []
+        with get_news_db() as conn:
+            for i in links:
+                cursor = conn.cursor()
+                cursor.execute("SELECT link, title, content FROM CNNHEALTHNEWS2 WHERE link = (?)", ([i]))
+                rows = cursor.fetchall()
+                prevs = [(str(row[0]), str(row[1]), str(row[2])) for row in rows]
+                preview_contents += prevs
+
+        return preview_contents
+    
+    except Exception as e:
+        print(f"Error in news_preview: {e}")
+        return ("", "", "")
+
+class Language(BaseModel):
+    query: str = Field(description="Translated query")
+    language: str = Field(description="Query's language")
+    
+def query_language(query):
+    system_prompt = f'''Your task is to determine what language the question is written in and translate it to english if it is not in English.
+    The output must be in this form: {{query: <translated query>, language: <query's language>}}
+    '''
+    
+    user_prompt = f'''
+    query: {query}
+    output:
+    '''
+    client = Together(api_key = API_KEY)
+    
+    response = client.chat.completions.create(
+       model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
+       response_format={
+            "type": "json_schema",
+            "schema": Language.model_json_schema(),
+       },
+       temperature = 0,
+       messages=[{
+           "role": "system",
+           "content": [
+               {"type": "text", "text":system_prompt}
+           ]
+       },
+                {
+           "role": "user",
+           "content": [
+               {"type": "text", "text":user_prompt},
+           ]
+       }])
+
+    return json.loads(response.choices[0].message.content)
+
+#API ENDPOINTS
+
+@app.post("/api/query", response_model=QueryResponse)
+def process_query(request: QueryRequest):
+    """Process user query and return comprehensive response"""
+    try:
+        # Step 1: Retrieve triplets
+        query = request.query
+        query = query_language(query)
+        
+        triplets_data, relations_data = retrieve_triplets(query['query'])
+        
+        # Step 2: Retrieve news
+        news_list, news_links = retrieve_news(query['query'])
+        
+        # Step 3: Extract information from triplets
+        triplet_info = extract_information_from_triplets(query['query'], triplets_data, relations_data)
+        
+        # Step 4: Extract information from news
+        news_info = extract_information_from_news(query['query'], news_list)
+        
+        # Step 5: Generate final answer
+        final_answer = extract_information(query['query'], triplet_info, news_info, query['language'])
+        
+        # Convert triplets to response format
+        triplets = [TripletData(head=t[0], relation=t[1], tail=t[2]) for t in triplets_data]
+        relations = [RelationDefinition(relation=r[0], definition=r[1]) for r in relations_data]
+        
+        # Convert news to response format with previews
+        news_prev = news_preview(news_links)
+        news_items = []
+        for url, title, content in news_prev:
+            preview = content[:300] + "..." if len(content) > 300 else content
+            news_items.append(NewsItem(
+                url=url,
+                content=content,
+                preview=preview,
+                title=title
+            ))
+        
+        # Create mini graph data for visualization
+        nodes_set = set()
+        edges = []
+        
+        for triplet in triplets_data:
+            head, relation, tail = triplet
+            nodes_set.add(head)
+            nodes_set.add(tail)
+            
+            # Find definition for this relation
+            definition = "No definition available"
+            for rel, def_text in relations_data:
+                if rel == relation:
+                    definition = def_text
+                    break
+            
+            edges.append(GraphEdge(
+                source=head,
+                target=tail,
+                relation=relation,
+                definition=definition
+            ))
+        
+        nodes = [GraphNode(id=node, label=node) for node in nodes_set]
+        graph_data = GraphData(nodes=nodes, edges=edges)
+        
+        return QueryResponse(
+            answer=final_answer,
+            triplets=triplets,
+            relations=relations,
+            news_items=news_items,
+            graph_data=graph_data
+        )
+        
+    except Exception as e:
+        print(f"Error in process_query: {e}")
+        raise HTTPException(status_code=500, detail=f"Query processing failed: {str(e)}")
+
+@app.get("/api/graph", response_model=GraphData)
+def get_graph_data(
+    search: Optional[str] = None,
+    triplets_db: sqlite3.Connection = Depends(get_triplets_connection),
+    definitions_db: sqlite3.Connection = Depends(get_definitions_connection)
+):
+    """Get complete graph data with nodes and edges."""
+    
+    try:
+        # Build dynamic query based on configuration
+        table = DATABASE_CONFIG["triplets_table"]
+        head_col = DATABASE_CONFIG["head_column"]
+        rel_col = DATABASE_CONFIG["relation_column"] 
+        tail_col = DATABASE_CONFIG["tail_column"]
+        
+        base_query = f"SELECT {head_col}, {rel_col}, {tail_col} FROM {table}"
+        params = []
+        
+        if search:
+            base_query += f" WHERE {head_col} LIKE ? OR {tail_col} LIKE ? OR {rel_col} LIKE ?"
+            search_term = f"%{search}%"
+            params = [search_term, search_term, search_term]
+        
+        base_query += " LIMIT 1000"
+        
+        # Get triplets
+        cursor = triplets_db.execute(base_query, params)
+        triplets = cursor.fetchall()
+
+        with get_definitions_db() as conn:
+            # Get definitions
+            def_table = DATABASE_CONFIG["definitions_table"]
+            def_col = DATABASE_CONFIG["definition_column"]
+            rel_col_def = DATABASE_CONFIG["relation_column"]
+            
+            def_cursor = conn.execute(f"SELECT {rel_col_def}, {def_col} FROM {def_table}")
+            definitions = {row[0]: row[1] for row in def_cursor.fetchall()}
+        
+        # Build nodes and edges
+        nodes_set = set()
+        edges = []
+        
+        for triple in triplets:
+            head = triple[0]
+            relation = triple[1]
+            tail = triple[2]
+            
+            # Add entities to nodes set
+            nodes_set.add(head)
+            nodes_set.add(tail)
+            
+            # Create edge with definition
+            edge = GraphEdge(
+                source=head,
+                target=tail,
+                relation=relation,
+                definition=definitions.get(relation, "No definition available")
+            )
+            edges.append(edge)
+        
+        # Convert nodes set to list of GraphNode objects
+        nodes = [GraphNode(id=node, label=node) for node in nodes_set]
+        
+        return GraphData(nodes=nodes, edges=edges)
+        
+    except Exception as e:
+        print(f"Error in get_graph_data: {e}")
+        raise HTTPException(status_code=500, detail=f"Database query failed: {str(e)}")
+
+if __name__ == "__main__":
+    print("Starting Knowledge Graph API...")
+    print(f"Triplets DB: {DATABASE_CONFIG['triplets_db']}")
+    print(f"Definitions DB: {DATABASE_CONFIG['definitions_db']}")
+    
+    import uvicorn
+    port = int(os.environ.get("PORT", 8000))
+    uvicorn.run(app, host="0.0.0.0", port=port)
+
+