@@ -108,6 +108,7 @@ def generate_with_langchain_docs(
108108 query_distribution : t .Optional [QueryDistribution ] = None ,
109109 run_config : t .Optional [RunConfig ] = None ,
110110 callbacks : t .Optional [Callbacks ] = None ,
111+ token_usage_parser : t .Optional [TokenUsageParser ] = None ,
111112 with_debugging_logs = False ,
112113 raise_exceptions : bool = True ,
113114 return_executor : bool = False ,
@@ -133,6 +134,9 @@ def generate_with_langchain_docs(
133134 Configuration for the generation run, by default None
134135 callbacks : Optional[Callbacks], optional
135136 Callbacks to use during generation, by default None
137+ token_usage_parser : Optional[TokenUsageParser], optional
138+ Parse the LLMResult object and return a TokenUsage object. This is used to
139+ calculate the cost of the generation process.
136140 with_debugging_logs : bool, optional
137141 Whether to include debug logs, by default False
138142 raise_exceptions : bool, optional
@@ -196,6 +200,7 @@ def generate_with_langchain_docs(
196200 query_distribution = query_distribution ,
197201 run_config = run_config ,
198202 callbacks = callbacks ,
203+ token_usage_parser = token_usage_parser ,
199204 with_debugging_logs = with_debugging_logs ,
200205 raise_exceptions = raise_exceptions ,
201206 return_executor = return_executor ,
@@ -211,6 +216,7 @@ def generate_with_llamaindex_docs(
211216 query_distribution : t .Optional [QueryDistribution ] = None ,
212217 run_config : t .Optional [RunConfig ] = None ,
213218 callbacks : t .Optional [Callbacks ] = None ,
219+ token_usage_parser : t .Optional [TokenUsageParser ] = None ,
214220 with_debugging_logs = False ,
215221 raise_exceptions : bool = True ,
216222 ):
@@ -274,6 +280,7 @@ def generate_with_llamaindex_docs(
274280 query_distribution = query_distribution ,
275281 run_config = run_config ,
276282 callbacks = callbacks ,
283+ token_usage_parser = token_usage_parser ,
277284 with_debugging_logs = with_debugging_logs ,
278285 raise_exceptions = raise_exceptions ,
279286 return_executor = False , # Default value for llamaindex_docs method
0 commit comments