1111import tiktoken
1212from jinja2 import Environment , FileSystemLoader , Template , TemplateNotFound
1313
14- from gitingest .schemas import ContextV1 , FileSystemNode , Source
14+ from gitingest .schemas import FileSystemNode , Source
1515from gitingest .schemas .filesystem import SEPARATOR , FileSystemNodeType
1616from gitingest .utils .compat_func import readlink
1717from gitingest .utils .logging_config import get_logger
@@ -136,15 +136,15 @@ def _format_token_count(text: str) -> str | None:
136136 return str (total_tokens )
137137
138138
139- def generate_digest (context : ContextV1 ) -> str :
140- """Generate a digest string from a ContextV1 object.
139+ def generate_digest (context : Source ) -> str :
140+ """Generate a digest string from a Source object.
141141
142- This is a convenience function that uses the DefaultFormatter to format a ContextV1 .
142+ This is a convenience function that uses the DefaultFormatter to format a Source .
143143
144144 Parameters
145145 ----------
146- context : ContextV1
147- The ContextV1 object containing sources and query information.
146+ context : Source
147+ The Source object containing sources and query information.
148148
149149 Returns
150150 -------
@@ -156,57 +156,103 @@ def generate_digest(context: ContextV1) -> str:
156156 return formatter .format (context , context .query )
157157
158158
159- class DefaultFormatter :
160- """Default formatter for rendering filesystem nodes using Jinja2 templates ."""
159+ class Formatter :
160+ """Base formatter class ."""
161161
162- def __init__ (self ) -> None :
162+ def __init__ (self , template_subdir : str ) -> None :
163163 self .separator = SEPARATOR
164- template_dir = Path (__file__ ).parent / "format" / "DefaultFormatter"
164+ template_dir = Path (__file__ ).parent / "format" / template_subdir
165165 self .env = Environment (loader = FileSystemLoader (template_dir ), autoescape = True )
166166
167167 def _get_template_for_node (self , node : Source ) -> Template :
168168 """Get template based on node class name."""
169169 template_name = f"{ node .__class__ .__name__ }
170170 return self .env .get_template (template_name )
171171
172+ 173+ class DefaultFormatter (Formatter ):
174+ """Default formatter for rendering filesystem nodes using Jinja2 templates."""
175+ 176+ def __init__ (self ) -> None :
177+ super ().__init__ ("DefaultFormatter" )
178+ 179+ def format (self , source : Source , query : IngestionQuery ) -> str :
180+ """Format a source with the given query."""
181+ if query is None :
182+ # Handle case where query is None (shouldn't happen in normal usage)
183+ raise ValueError ("ContextV1 must have a valid query object" )
184+ 185+ # Calculate and set token count for ContextV1
186+ if hasattr (source , '_token_count' ):
187+ token_count = self ._calculate_token_count (source )
188+ source ._token_count = token_count
189+ # Also set token count in the extra dict
190+ source .extra ["token_count" ] = token_count
191+ 192+ try :
193+ return self ._format_node (source , query )
194+ except Exception as e :
195+ # Log the error for debugging
196+ import logging
197+ logging .error (f"Error in DefaultFormatter: { e } )
198+ raise
199+ 200+ def _calculate_token_count (self , source : Source ) -> str :
201+ """Calculate token count for the entire source."""
202+ # Gather all content from the source
203+ content = self ._gather_all_content (source )
204+ return _format_token_count (content ) or "Unknown"
205+ 206+ def _gather_all_content (self , node : Source ) -> str :
207+ """Recursively gather all content from the source tree."""
208+ content_parts = []
209+ 210+ # Add content from the current node
211+ if hasattr (node , 'content' ):
212+ content_parts .append (node .content )
213+ 214+ # Add content from all sources if it's a ContextV1
215+ if hasattr (node , 'sources' ):
216+ for source in node .sources :
217+ content_parts .append (self ._gather_all_content (source ))
218+ 219+ # Add content from children if it's a directory
220+ if hasattr (node , 'children' ):
221+ for child in node .children :
222+ content_parts .append (self ._gather_all_content (child ))
223+ 224+ return "\n " .join (filter (None , content_parts ))
225+ 172226 @singledispatchmethod
173- def format (self , node : Source , query : IngestionQuery ) -> str :
227+ def _format_node (self , node : Source , query : IngestionQuery ) -> str :
174228 """Dynamically format any node type based on available templates."""
175229 try :
176230 template = self ._get_template_for_node (node )
177231 # Provide common template variables
178232 context_vars = {
179- "node " : node ,
233+ "source " : node ,
180234 "query" : query ,
181235 "formatter" : self ,
182236 "SEPARATOR" : SEPARATOR ,
183237 }
184- # Special handling for ContextV1 objects
185- if isinstance (node , ContextV1 ):
186- context_vars ["context" ] = node
187- # Use ContextV1 for backward compatibility
188- template = self .env .get_template ("ContextV1.j2" )
189238
190239 return template .render (** context_vars )
191240 except TemplateNotFound :
192241 # Fallback: return content if available, otherwise empty string
193242 return f"{ getattr (node , 'content' , '' )}
194243
195244
196- class DebugFormatter :
245+ class DebugFormatter ( Formatter ) :
197246 """Debug formatter that shows detailed information about filesystem nodes."""
198247
199248 def __init__ (self ) -> None :
200- self .separator = SEPARATOR
201- template_dir = Path (__file__ ).parent / "format" / "DebugFormatter"
202- self .env = Environment (loader = FileSystemLoader (template_dir ), autoescape = True )
249+ super ().__init__ ("DebugFormatter" )
203250
204251 def _get_template_for_node (self , node : Source ) -> Template :
205252 """Get template based on node class name."""
206253 template_name = f"{ node .__class__ .__name__ }
207254 return self .env .get_template (template_name )
208255
209- @singledispatchmethod
210256 def format (self , node : Source , query : IngestionQuery ) -> str :
211257 """Dynamically format any node type with debug information."""
212258 try :
@@ -254,37 +300,31 @@ def _raise_no_dataclass_fields() -> None:
254300 return f"DEBUG: { node .__class__ .__name__ }
255301
256302
257- class SummaryFormatter :
303+ class SummaryFormatter ( Formatter ) :
258304 """Dedicated formatter for generating summaries of filesystem nodes."""
259305
260306 def __init__ (self ) -> None :
261- template_dir = Path (__file__ ).parent / "format" / "SummaryFormatter"
262- self .env = Environment (loader = FileSystemLoader (template_dir ), autoescape = True )
307+ super ().__init__ ("SummaryFormatter" )
263308
264- def _get_template_for_node (self , node : Source ) -> Template :
265- """Get template based on node class name."""
266- template_name = f"{ node .__class__ .__name__ }
267- return self .env .get_template (template_name )
309+ def format (self , source : Source , query : IngestionQuery ) -> str :
310+ """Generate the summary output."""
311+ if query is None :
312+ # Handle case where query is None (shouldn't happen in normal usage)
313+ raise ValueError ("ContextV1 must have a valid query object" )
314+ return self .summary (source , query )
268315
269316 @singledispatchmethod
270317 def summary (self , node : Source , query : IngestionQuery ) -> str :
271318 """Dynamically generate summary for any node type based on available templates."""
272319 try :
273320 # Provide common template variables
274321 context_vars = {
275- "node " : node ,
322+ "source " : node ,
276323 "query" : query ,
277324 "formatter" : self ,
278325 }
279326
280- # Special handling for ContextV1 objects
281- if isinstance (node , ContextV1 ):
282- context_vars ["context" ] = node
283- # Use ContextV1 for backward compatibility
284- template = self .env .get_template ("ContextV1.j2" )
285- else :
286- template = self ._get_template_for_node (node )
287- 327+ template = self ._get_template_for_node (node )
288328 return template .render (** context_vars )
289329 except TemplateNotFound :
290330 # Fallback: return name if available
0 commit comments