allow for saving logs in parser

guidance-ai · Jul 31, 2024 · bb98717 · bb98717
1 parent 94ccd77
commit bb98717
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 10 deletions.
diff --git a/parser/src/earley/from_guidance.rs b/parser/src/earley/from_guidance.rs
@@ -1,3 +1,4 @@
+use std::fmt::Write;
 use std::{sync::Arc, vec};
 
 use super::{grammar::SymbolProps, lexerspec::LexerSpec, CGrammar, Grammar};
@@ -199,7 +200,10 @@ fn grammar_from_json(input: GrammarWithLexer) -> Result<(LexerSpec, Grammar)> {
     Ok((lexer_spec, grm))
 }
 
-pub fn grammars_from_json(input: TopLevelGrammar, print_out: bool) -> Result<Vec<Arc<CGrammar>>> {
+pub fn grammars_from_json(
+    input: TopLevelGrammar,
+    print_out: bool,
+) -> Result<(Vec<Arc<CGrammar>>, String)> {
     let grammars = input
         .grammars
         .into_iter()
@@ -210,21 +214,24 @@ pub fn grammars_from_json(input: TopLevelGrammar, print_out: bool) -> Result<Vec
         g.validate_grammar_refs(&grammars)?;
     }
 
-    Ok(grammars
+    let mut out = String::new();
+    let grammars = grammars
         .into_iter()
         .enumerate()
         .map(|(idx, (lex, mut grm))| {
             if print_out {
-                eprintln!("\nGrammar #{}:\n{:?}\n{:?}", idx, lex, grm);
+                writeln!(out, "Grammar #{}:\n{:?}\n{:?}\n", idx, lex, grm).unwrap();
             }
 
             grm = grm.optimize();
 
             if print_out {
-                eprintln!("  == Optimize ==>\n{:?}", grm);
+                write!(out, "  == Optimize ==>\n{:?}", grm).unwrap();
             }
 
             Arc::new(grm.compile(lex))
         })
-        .collect::<Vec<_>>())
+        .collect::<Vec<_>>();
+
+    Ok((grammars, out))
 }
diff --git a/parser/src/tokenparser.rs b/parser/src/tokenparser.rs
@@ -1,3 +1,4 @@
+use std::fmt::Write;
 use std::sync::Arc;
 
 use crate::{
@@ -11,16 +12,25 @@ use toktrie::{InferenceCapabilities, SimpleVob, StepArg, StepResult, TokenId, To
 macro_rules! infoln {
     ($s:expr, $($arg:tt)*) => {
         if $s.log_level >= 2 {
-            eprintln!($($arg)*);
+            if $s.save_logs {
+                writeln!($s.logs, $($arg)*).unwrap();
+            } else {
+                eprintln!($($arg)*);
+            }
         }
     };
 }
 
 macro_rules! warn {
     ($s:expr, $($arg:tt)*) => {
         if $s.log_level >= 1 {
-            eprint!("Warning: ");
-            eprintln!($($arg)*);
+            if $s.save_logs {
+                $s.logs.push_str("Warning: ");
+                writeln!($s.logs, $($arg)*).unwrap();
+            } else {
+                eprint!("Warning: ");
+                eprintln!($($arg)*);
+            }
         }
     };
 }
@@ -32,6 +42,8 @@ pub struct TokenParser {
     pub log_level: isize,
     pub mid_process_start_time: std::time::Instant,
     pub inference_caps: InferenceCapabilities,
+    save_logs: bool,
+    logs: String,
     pending_bogus_backtrack: u32,
     // sampling any of these will pop the parser stack:
     pop_tokens: Option<SimpleVob>,
@@ -72,19 +84,22 @@ impl TokenParser {
         token_env: Arc<dyn TokenizerEnv + Sync>,
         buf: TopLevelGrammar,
         log_level: isize,
+        save_logs: bool,
         inference_caps: InferenceCapabilities,
     ) -> Result<Self> {
         let mid_process_start_time = std::time::Instant::now();
         let test_trace = buf.test_trace;
         let max_tokens = buf.max_tokens.unwrap_or(usize::MAX);
-        let compiled_grammars = grammars_from_json(buf, log_level >= 2)?;
+        let (compiled_grammars, grammar_log) = grammars_from_json(buf, log_level >= 2)?;
         let parser = Parser::new(
             Arc::clone(&compiled_grammars[0]),
             GenGrammarOptions::default(),
         )?;
 
         Ok(TokenParser {
             log_level,
+            logs: grammar_log,
+            save_logs,
             test_trace,
             token_env,
             inference_caps,
@@ -107,6 +122,10 @@ impl TokenParser {
         })
     }
 
+    pub fn get_and_clear_logs(&mut self) -> String {
+        std::mem::replace(&mut self.logs, String::new())
+    }
+
     pub fn stop_reason(&self) -> StopReason {
         self.stop_reason
     }
@@ -209,7 +228,7 @@ impl TokenParser {
         self.parser_stack.is_empty()
     }
 
-    fn test_trace_json(&self, j: &serde_json::Value) {
+    fn test_trace_json(&mut self, j: &serde_json::Value) {
         if self.test_trace {
             infoln!(self, "TEST: {}", serde_json::to_string(j).unwrap());
         }