Citations
This commit is contained in:
Yiorgis Gozadinos
parent
f8964f6efc
commit
9f1d82daa8
3 changed files with 228 additions and 10 deletions
|
|
@ -13,10 +13,24 @@ if TYPE_CHECKING:
|
|||
from haiku.rag.graph.agui.emitter import AGUIEmitter
|
||||
|
||||
|
||||
class CitationInfo(BaseModel):
|
||||
"""Citation info for frontend display."""
|
||||
|
||||
index: int
|
||||
document_id: str
|
||||
chunk_id: str
|
||||
document_uri: str
|
||||
document_title: str | None = None
|
||||
page_numbers: list[int] = []
|
||||
headings: list[str] | None = None
|
||||
content: str
|
||||
|
||||
|
||||
class ChatSessionState(BaseModel):
|
||||
"""State shared between frontend and agent via AG-UI."""
|
||||
|
||||
session_id: str = ""
|
||||
citations: list[CitationInfo] = []
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -35,13 +49,12 @@ You have access to a knowledge base of documents. Use your tools to search and a
|
|||
|
||||
CRITICAL RULES:
|
||||
1. For greetings or casual chat: respond directly WITHOUT using any tools
|
||||
2. For substantive questions requiring information: use the search or ask tools
|
||||
2. For questions: ALWAYS use the "ask" tool - it provides answers with proper citations
|
||||
3. NEVER make up information - always use tools to get facts from the knowledge base
|
||||
4. When citing sources, reference the chunk IDs from search results
|
||||
|
||||
How to decide which tool to use:
|
||||
- "search" - When you need to find relevant documents or explore what's in the knowledge base
|
||||
- "ask" - When you have a specific question that needs a direct answer with citations
|
||||
- "ask" - DEFAULT CHOICE for any question. Use this for questions like "What is X?", "How does Y work?", "Explain Z", etc. Returns answers with citations.
|
||||
- "search" - ONLY use when explicitly exploring/browsing the knowledge base, or when the user asks to "search for" or "find" something without needing an answer.
|
||||
|
||||
Be friendly and conversational. When you use tools, summarize the key findings for the user."""
|
||||
|
||||
|
|
@ -111,13 +124,29 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
|
||||
answer, citations = await ctx.deps.client.ask(question, filter=document_filter)
|
||||
|
||||
# Format answer with citations
|
||||
if citations:
|
||||
citation_list = "\n".join(
|
||||
f" [{i + 1}] {c.document_uri or c.document_title or 'Unknown'} (chunk: {c.chunk_id})"
|
||||
# Emit citations via AG-UI state for frontend rendering
|
||||
if citations and ctx.deps.agui_emitter:
|
||||
citation_infos = [
|
||||
CitationInfo(
|
||||
index=i + 1,
|
||||
document_id=c.document_id,
|
||||
chunk_id=c.chunk_id,
|
||||
document_uri=c.document_uri,
|
||||
document_title=c.document_title,
|
||||
page_numbers=c.page_numbers,
|
||||
headings=c.headings,
|
||||
content=c.content,
|
||||
)
|
||||
for i, c in enumerate(citations)
|
||||
]
|
||||
ctx.deps.agui_emitter.update_state(
|
||||
ChatSessionState(citations=citation_infos)
|
||||
)
|
||||
return f"{answer}\n\nSources:\n{citation_list}"
|
||||
|
||||
# Format answer with citation references
|
||||
if citations:
|
||||
citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citations)))
|
||||
return f"{answer}\n\nSources: {citation_refs}"
|
||||
|
||||
return answer
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,29 @@
|
|||
"use client";
|
||||
|
||||
import { CopilotKit, useCoAgent } from "@copilotkit/react-core";
|
||||
import {
|
||||
CopilotKit,
|
||||
useCoAgent,
|
||||
useCoAgentStateRender,
|
||||
} from "@copilotkit/react-core";
|
||||
import { CopilotChat } from "@copilotkit/react-ui";
|
||||
import "@copilotkit/react-ui/styles.css";
|
||||
import CitationBlock from "./CitationBlock";
|
||||
import DbInfo from "./DbInfo";
|
||||
|
||||
interface Citation {
|
||||
index: number;
|
||||
document_id: string;
|
||||
chunk_id: string;
|
||||
document_uri: string;
|
||||
document_title: string | null;
|
||||
page_numbers: number[];
|
||||
headings: string[] | null;
|
||||
content: string;
|
||||
}
|
||||
|
||||
interface ChatSessionState {
|
||||
session_id: string;
|
||||
citations: Citation[];
|
||||
}
|
||||
|
||||
function ChatContent() {
|
||||
|
|
@ -14,6 +31,17 @@ function ChatContent() {
|
|||
name: "chat_agent",
|
||||
initialState: {
|
||||
session_id: "",
|
||||
citations: [],
|
||||
},
|
||||
});
|
||||
|
||||
useCoAgentStateRender<ChatSessionState>({
|
||||
name: "chat_agent",
|
||||
render: ({ state }) => {
|
||||
if (state.citations && state.citations.length > 0) {
|
||||
return <CitationBlock citations={state.citations} />;
|
||||
}
|
||||
return null;
|
||||
},
|
||||
});
|
||||
|
||||
|
|
@ -41,6 +69,12 @@ function ChatContent() {
|
|||
.chat-content {
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
}
|
||||
.chat-content > * {
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
}
|
||||
`}</style>
|
||||
<div className="chat-wrapper">
|
||||
|
|
|
|||
155
app/frontend/components/CitationBlock.tsx
Normal file
155
app/frontend/components/CitationBlock.tsx
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
"use client";
|
||||
|
||||
import { useState } from "react";
|
||||
|
||||
interface Citation {
|
||||
index: number;
|
||||
document_id: string;
|
||||
chunk_id: string;
|
||||
document_uri: string;
|
||||
document_title: string | null;
|
||||
page_numbers: number[];
|
||||
headings: string[] | null;
|
||||
content: string;
|
||||
}
|
||||
|
||||
interface CitationBlockProps {
|
||||
citations: Citation[];
|
||||
}
|
||||
|
||||
function CitationItem({ citation }: { citation: Citation }) {
|
||||
const [expanded, setExpanded] = useState(false);
|
||||
|
||||
const title = citation.document_title || citation.document_uri || "Unknown";
|
||||
const pageInfo =
|
||||
citation.page_numbers.length > 0
|
||||
? `p. ${citation.page_numbers.join(", ")}`
|
||||
: null;
|
||||
|
||||
return (
|
||||
<div className="citation-item">
|
||||
<button
|
||||
type="button"
|
||||
className="citation-header"
|
||||
onClick={() => setExpanded(!expanded)}
|
||||
>
|
||||
<span className="citation-index">[{citation.index}]</span>
|
||||
<span className="citation-title">{title}</span>
|
||||
{pageInfo && <span className="citation-page">{pageInfo}</span>}
|
||||
<span className={`citation-chevron ${expanded ? "expanded" : ""}`}>
|
||||
{expanded ? "▼" : "▶"}
|
||||
</span>
|
||||
</button>
|
||||
{expanded && (
|
||||
<div className="citation-content">
|
||||
{citation.headings && citation.headings.length > 0 && (
|
||||
<div className="citation-headings">
|
||||
{citation.headings.join(" › ")}
|
||||
</div>
|
||||
)}
|
||||
<div className="citation-text">{citation.content}</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
export default function CitationBlock({ citations }: CitationBlockProps) {
|
||||
if (!citations || citations.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<style>{`
|
||||
.citation-block {
|
||||
margin-top: 0.75rem;
|
||||
border: 1px solid #e2e8f0;
|
||||
border-radius: 8px;
|
||||
overflow: hidden;
|
||||
font-size: 0.875rem;
|
||||
}
|
||||
.citation-block-header {
|
||||
background: #f8fafc;
|
||||
padding: 0.5rem 0.75rem;
|
||||
font-weight: 500;
|
||||
color: #475569;
|
||||
border-bottom: 1px solid #e2e8f0;
|
||||
}
|
||||
.citation-item {
|
||||
border-bottom: 1px solid #f1f5f9;
|
||||
}
|
||||
.citation-item:last-child {
|
||||
border-bottom: none;
|
||||
}
|
||||
.citation-header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.5rem;
|
||||
width: 100%;
|
||||
padding: 0.5rem 0.75rem;
|
||||
background: none;
|
||||
border: none;
|
||||
cursor: pointer;
|
||||
text-align: left;
|
||||
transition: background 0.15s;
|
||||
}
|
||||
.citation-header:hover {
|
||||
background: #f8fafc;
|
||||
}
|
||||
.citation-index {
|
||||
font-weight: 600;
|
||||
color: #3b82f6;
|
||||
flex-shrink: 0;
|
||||
}
|
||||
.citation-title {
|
||||
flex: 1;
|
||||
color: #334155;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
}
|
||||
.citation-page {
|
||||
color: #94a3b8;
|
||||
font-size: 0.75rem;
|
||||
flex-shrink: 0;
|
||||
}
|
||||
.citation-chevron {
|
||||
color: #94a3b8;
|
||||
font-size: 0.625rem;
|
||||
flex-shrink: 0;
|
||||
transition: transform 0.15s;
|
||||
}
|
||||
.citation-chevron.expanded {
|
||||
transform: rotate(0deg);
|
||||
}
|
||||
.citation-content {
|
||||
padding: 0.75rem;
|
||||
background: #fafafa;
|
||||
border-top: 1px solid #f1f5f9;
|
||||
}
|
||||
.citation-headings {
|
||||
font-size: 0.75rem;
|
||||
color: #64748b;
|
||||
margin-bottom: 0.5rem;
|
||||
font-style: italic;
|
||||
}
|
||||
.citation-text {
|
||||
color: #475569;
|
||||
line-height: 1.5;
|
||||
white-space: pre-wrap;
|
||||
max-height: 200px;
|
||||
overflow-y: auto;
|
||||
}
|
||||
`}</style>
|
||||
<div className="citation-block">
|
||||
<div className="citation-block-header">
|
||||
Sources ({citations.length})
|
||||
</div>
|
||||
{citations.map((citation) => (
|
||||
<CitationItem key={citation.chunk_id} citation={citation} />
|
||||
))}
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
}
|
||||
Loading…
Reference in a new issue