traits.rs 2.46 KB
Newer Older
1
2
use async_trait::async_trait;

3
4
5
6
7
8
9
10
use crate::{
    protocols::common::Tool,
    tool_parser::{
        errors::ParserResult,
        types::{StreamingParseResult, ToolCall},
    },
};

11
12
13
14
/// Core trait for all tool parsers
#[async_trait]
pub trait ToolParser: Send + Sync {
    /// Parse complete tool calls from final output
15
    /// Returns (remaining_normal_text, tool_calls) tuple
16
    async fn parse_complete(&self, output: &str) -> ParserResult<(String, Vec<ToolCall>)>;
17
18

    /// Parse tool calls from model output (streaming)
19
20
21
22
23
    /// Parsers now maintain internal state, so self is mutable
    ///
    /// # Arguments
    /// * `chunk` - New text chunk from model output
    /// * `tools` - List of available tools for validation
24
    async fn parse_incremental(
25
        &mut self,
26
        chunk: &str,
27
        tools: &[Tool],
28
    ) -> ParserResult<StreamingParseResult>;
29
30

    /// Check if text contains tool calls in this parser's format
31
    fn has_tool_markers(&self, text: &str) -> bool;
32
33
34
35
36
37

    /// Optionally expose a token-aware parser implementation.
    /// Default returns `None`, meaning the parser only supports text input.
    fn as_token_parser(&self) -> Option<&dyn TokenToolParser> {
        None
    }
38
39
40
41
42
43

    /// Get unstreamed tool call arguments
    /// Returns tool call items for arguments that have been parsed but not yet streamed
    fn get_unstreamed_tool_args(&self) -> Option<Vec<crate::tool_parser::types::ToolCallItem>> {
        None
    }
44
45
46
47
48
49

    /// Reset the parser state for reuse across requests.
    /// This should clear all buffers and reset state to initial values.
    fn reset(&mut self) {
        // Default no-op implementation
    }
50
51
52
53
54
}

/// Trait for partial JSON parsing
pub trait PartialJsonParser: Send + Sync {
    /// Parse potentially incomplete JSON
55
    fn parse(&self, input: &str) -> ParserResult<(serde_json::Value, usize)>;
56
57
58
59
60
61
62

    /// Check if JSON is complete
    fn is_complete(&self, input: &str) -> bool;

    /// Get the maximum parsing depth
    fn max_depth(&self) -> usize;
}
63
64
65
66

#[async_trait]
pub trait TokenToolParser: ToolParser {
    /// Parse complete tool calls when provided with raw token IDs.
67
    async fn parse_complete_tokens(&self, tokens: &[u32]) -> ParserResult<(String, Vec<ToolCall>)>;
68
69

    /// Streaming parser entrypoint for token chunks.
70
    /// Parsers maintain internal state, so self is mutable
71
    async fn parse_incremental_tokens(
72
        &mut self,
73
        tokens: &[u32],
74
        tools: &[Tool],
75
    ) -> ParserResult<StreamingParseResult>;
76
}