Work with inline citations in Knowledge Graph responses

from writerai import Writer
import re

# Initialize the Writer client
client = Writer()

# Query Knowledge Graph with inline citations
response = client.graphs.question(
    graph_ids=["<GRAPH_ID>"],
    question="What are the key features of our product?",
    query_config={
        "inline_citations": True
    }
)

def extract_citations_with_sources(answer_text, references):
    '''Extract citations from answer text and link them to source material.
    
    This function parses inline citations from Knowledge Graph responses and correlates
    them with their source material from the references object.
    
    Args:
        answer_text (str): The response text containing inline citations
        references (object): The references object containing files and web sources
        
    Returns:
        list: Array of citation objects with source text and metadata
    '''
    # Use regex to find all citations in the format [display_text](identifier)
    # This pattern matches markdown-style links like [filename](citation-id)
    citation_pattern = r'\[([^\]]+)\]\(([^)]+)\)'
    citations = re.findall(citation_pattern, answer_text)
    
    citation_data = []
    
    # Process each citation found in the answer text
    for display_text, identifier in citations:
        # Initialize citation info structure
        citation_info = {
            "display_text": display_text,  # The text shown in the citation (e.g., filename)
            "identifier": identifier,  # The unique citation ID
            "source_text": None,  # The actual source text
            "source_type": None   # The source type: "file" or "web"
        }
        
        # Search through file references to find matching citation
        if references and hasattr(references, 'files') and references.files:
            for file_ref in references.files:
                # Match citation ID with the 'cite' field in file references
                if hasattr(file_ref, 'cite') and file_ref.cite == identifier:
                    citation_info["source_text"] = file_ref.text  # Extract source text snippet
                    citation_info["source_type"] = "file"  # Mark as file source
                    citation_info["file_id"] = file_ref.file_id  # Store file ID for reference
                    citation_info["page"] = file_ref.page  # Store page number if available
                    break
        
        # Search through web references to find matching citation
        if references and hasattr(references, 'web') and references.web:
            for web_ref in references.web:
                # Match citation ID with the 'url' field in web references
                if hasattr(web_ref, 'url') and web_ref.url == identifier:
                    citation_info["source_text"] = web_ref.text  # Extract source text snippet
                    citation_info["source_type"] = "web"  # Mark as web source
                    citation_info["url"] = web_ref.url  # Store URL for reference
                    citation_info["title"] = web_ref.title  # Store page title
                    break
        
        citation_data.append(citation_info)
    
    return citation_data

# Extract citations with their source text
# This processes the response and creates a structured list of citations with their source material
citations = extract_citations_with_sources(response.answer, response.references)

# Display citations with source material
# Loop through each citation and print its details along with the source text
for citation in citations:
    print(f"Citation: [{citation['display_text']}]({citation['identifier']})")
    print(f"Source type: {citation['source_type']}")
    print(f"Source text: {citation['source_text']}")
    
    # Show additional metadata based on source type
    if citation['source_type'] == 'file':
        print(f"File ID: {citation['file_id']}, Page: {citation['page']}")
    elif citation['source_type'] == 'web':
        print(f"URL: {citation['url']}, Title: {citation['title']}")
    
    print("-" * 50)  # Separator between citations

from writerai import Writer

# Initialize the Writer client
client = Writer()

# Configure Knowledge Graph tool with inline citations
tools = [{
    "type": "graph",
    "function": {
        "description": "Search company knowledge base",
        "graph_ids": ["<GRAPH_ID>"],
        "subqueries": True,
        "query_config": {
            "inline_citations": True
        }
    }
}]

messages = [{'role': 'user', 'content': 'What are the key features of our product?'}]

# Stream the response
stream = client.chat.chat(
    model="palmyra-x5",
    messages=messages,
    tools=tools,
    tool_choice="auto",
    stream=True
)

# Collect streaming content and references
references = None

for chunk in stream:
    if chunk.choices[0].delta.content is not None:
        # Stream content progressively
        content = chunk.choices[0].delta.content
        print(content, end='', flush=True)
    
    # Check for graph data in the message
    if hasattr(chunk.choices[0], 'message') and chunk.choices[0].message:
        if hasattr(chunk.choices[0].message, 'graph_data') and chunk.choices[0].message.graph_data:
            # References are in the graph_data.references object
            references = chunk.choices[0].message.graph_data.references

print("\nFile references:", references.files)
print("\nWeb references:", references.web)