Skip to content

Commit

Permalink
feat: add robustness to get_source_code_chunks
Browse files Browse the repository at this point in the history
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
  • Loading branch information
lxobr and coderabbitai[bot] authored Dec 23, 2024
1 parent 35071b5 commit 68a9d27
Showing 1 changed file with 16 additions and 9 deletions.
25 changes: 16 additions & 9 deletions cognee/tasks/repo_processor/get_source_code_chunks.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,12 +146,19 @@ async def get_source_code_chunks(data_points: list[DataPoint], embedding_model="
"""Processes code graph datapoints, create SourceCodeChink datapoints."""
# TODO: Add support for other embedding models, with max_token mapping
for data_point in data_points:
yield data_point
if not isinstance(data_point, CodeFile):
continue
if not data_point.contains:
continue
for code_part in data_point.contains:
yield code_part
for source_code_chunk in get_source_code_chunks_from_code_part(code_part, model_name=embedding_model):
yield source_code_chunk
try:
yield data_point
if not isinstance(data_point, CodeFile):
continue
if not data_point.contains:
logger.warning(f"CodeFile {data_point.id} contains no code parts")
continue
for code_part in data_point.contains:
try:
yield code_part
for source_code_chunk in get_source_code_chunks_from_code_part(code_part, model_name=embedding_model):
yield source_code_chunk
except Exception as e:
logger.error(f"Error processing code part: {e}")
except Exception as e:
logger.error(f"Error processing data point: {e}")

0 comments on commit 68a9d27

Please sign in to comment.