gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

ruff changes

Browse source
This commit is contained in:
rajeevrajeshuni 2025-11-30 14:20:25 +05:30
parent e4a2188820
commit 37e862c053
13 changed files with 120 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/cli/commands/delete_command.py
View file

@ -109,4 +109,4 @@ Be careful with deletion operations as they are irreversible.
except Exception as e:
if isinstance(e, CliCommandInnerException):
raise CliCommandException(str(e), error_code=1) from e
raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1) from e
raise CliCommandException(f"Error deleting data: {str(e)}", error_code=1) from e

1
cognee/cli/commands/tui_command.py
View file

@ -27,6 +27,7 @@ class TuiCommand(SupportsCliCommand):
def execute(self, args: argparse.Namespace) -> None:
try:
class CogneeTUI(App):
"""Main TUI application for cognee."""

35
cognee/cli/tui/add_screen.py
View file

@ -17,12 +17,15 @@ class AddTUIScreen(BaseTUIScreen):
Binding("ctrl+v", "paste", "Paste", show=False),
]
CSS = BaseTUIScreen.CSS + """
CSS = (
BaseTUIScreen.CSS
+ """
#data-input {
height: 8;
min-height: 8;
}
"""
)
def __init__(self):
super().__init__()
@ -33,25 +36,21 @@ class AddTUIScreen(BaseTUIScreen):
with Container(classes="tui-title-wrapper"):
yield Static("📥 Add Data to Cognee", classes="tui-title-bordered")
with Vertical(classes="tui-form"):
yield Label("Data (text, file path (/path/to/doc), URL, or S3 path (s3://bucket)):", classes="tui-label-spaced")
yield Label(
"Data (text, file path (/path/to/doc), URL, or S3 path (s3://bucket)):",
classes="tui-label-spaced",
)
yield TextArea(
"",
id="data-input",
)
yield Label("Dataset Name:", classes="tui-label-spaced")
yield Input(
placeholder="main_dataset",
value="main_dataset",
id="dataset-input"
)
yield Input(placeholder="main_dataset", value="main_dataset", id="dataset-input")
yield Static("", classes="tui-status")
def compose_footer(self) -> ComposeResult:
yield Static(
"Ctrl+S: Add • Esc: Back • q: Quit",
classes="tui-footer"
)
yield Static("Ctrl+S: Add • Esc: Back • q: Quit", classes="tui-footer")
def on_mount(self) -> None:
"""Focus the data input on mount."""
@ -91,7 +90,7 @@ class AddTUIScreen(BaseTUIScreen):
self.is_processing = True
status.update("[yellow]⏳ Processing...[/yellow]")
# Disable inputs during processing
data_input.disabled = True
dataset_input.disabled = True
@ -102,21 +101,21 @@ class AddTUIScreen(BaseTUIScreen):
async def _add_data_async(self, data: str, dataset_name: str) -> None:
"""Async function to add data to cognee."""
status = self.query_one(".tui-status", Static)
try:
import cognee
await cognee.add(data=data, dataset_name=dataset_name)
status.update(f"[green]✓ Successfully added data to dataset '{dataset_name}'[/green]")
# Clear the data input after successful add
data_input = self.query_one("#data-input", TextArea)
data_input.clear()
except Exception as e:
status.update(f"[red]✗ Failed to add data: {str(e)}[/red]")
finally:
# Re-enable inputs
self.is_processing = False

5
cognee/cli/tui/base_screen.py
View file

@ -10,7 +10,9 @@ class BaseTUIScreen(Screen):
"""Base screen class with constant header for all TUI screens."""
# Subclasses should override this CSS and add their own styles
CSS = COMMON_STYLES + """
CSS = (
COMMON_STYLES
+ """
#header {
dock: top;
background: $boost;
@ -21,6 +23,7 @@ class BaseTUIScreen(Screen):
padding: 1;
}
"""
)
def compose_header(self) -> ComposeResult:
"""Compose the constant header widget."""

50
cognee/cli/tui/cognify_screen.py
View file

@ -16,7 +16,9 @@ class CognifyTUIScreen(BaseTUIScreen):
Binding("ctrl+s", "submit", "Submit"),
]
CSS = BaseTUIScreen.CSS + """
CSS = (
BaseTUIScreen.CSS
+ """
Checkbox {
margin-top: 1;
margin-bottom: 1;
@ -32,6 +34,7 @@ class CognifyTUIScreen(BaseTUIScreen):
height: 1;
}
"""
)
def __init__(self):
super().__init__()
@ -42,26 +45,23 @@ class CognifyTUIScreen(BaseTUIScreen):
with Container(classes="tui-title-wrapper"):
yield Static("⚡ Cognify Data", classes="tui-title-bordered")
with Vertical(classes="tui-form"):
yield Label("Dataset Name (optional, leave empty for all):", classes="tui-label-spaced")
yield Input(
placeholder="Leave empty to process all datasets",
value="",
id="dataset-input"
yield Label(
"Dataset Name (optional, leave empty for all):", classes="tui-label-spaced"
)
yield Input(
placeholder="Leave empty to process all datasets", value="", id="dataset-input"
)
yield Label("Chunker Type:", classes="tui-label-spaced")
with RadioSet(id="chunker-radio"):
for chunker in CHUNKER_CHOICES:
yield RadioButton(chunker, value=(chunker == "TextChunker"))
yield Checkbox("Run in background", id="background-checkbox")
yield Static("", classes="tui-status")
def compose_footer(self) -> ComposeResult:
yield Static(
"Ctrl+S: Start • Esc: Back • q: Quit",
classes="tui-footer"
)
yield Static("Ctrl+S: Start • Esc: Back • q: Quit", classes="tui-footer")
def on_mount(self) -> None:
"""Focus the dataset input on mount."""
@ -90,12 +90,16 @@ class CognifyTUIScreen(BaseTUIScreen):
status = self.query_one(".tui-status", Static)
dataset_name = dataset_input.value.strip() or None
chunker_type = str(chunker_radio.pressed_button.label) if chunker_radio.pressed_button else "TextChunker"
chunker_type = (
str(chunker_radio.pressed_button.label)
if chunker_radio.pressed_button
else "TextChunker"
)
run_background = background_checkbox.value
self.is_processing = True
status.update("[yellow]⏳ Starting cognification...[/yellow]")
# Disable inputs during processing
dataset_input.disabled = True
chunker_radio.disabled = True
@ -104,10 +108,13 @@ class CognifyTUIScreen(BaseTUIScreen):
# Run async cognify operation
asyncio.create_task(self._cognify_async(dataset_name, chunker_type, run_background))
async def _cognify_async(self, dataset_name: str | None, chunker_type: str, run_background: bool) -> None:
async def _cognify_async(
self, dataset_name: str | None, chunker_type: str, run_background: bool
) -> None:
"""Async function to cognify data."""
status = self.query_one(".tui-status", Static)
from cognee.modules.chunking.TextChunker import TextChunker
try:
# Get chunker class
chunker_class = TextChunker
@ -119,7 +126,9 @@ class CognifyTUIScreen(BaseTUIScreen):
if LangchainChunker is not None:
chunker_class = LangchainChunker
else:
status.update("[yellow]⚠ LangchainChunker not available, using TextChunker[/yellow]")
status.update(
"[yellow]⚠ LangchainChunker not available, using TextChunker[/yellow]"
)
elif chunker_type == "CsvChunker":
try:
from cognee.modules.chunking.CsvChunker import CsvChunker
@ -129,24 +138,25 @@ class CognifyTUIScreen(BaseTUIScreen):
chunker_class = CsvChunker
else:
status.update("[yellow]⚠ CsvChunker not available, using TextChunker[/yellow]")
# Prepare datasets parameter
datasets = [dataset_name] if dataset_name else None
import cognee
await cognee.cognify(
datasets=datasets,
chunker=chunker_class,
run_in_background=run_background,
)
if run_background:
status.update("[green]✓ Cognification started in background![/green]")
else:
status.update("[green]✓ Cognification completed successfully![/green]")
except Exception as e:
status.update(f"[red]✗ Failed to cognify: {str(e)}[/red]")
finally:
# Re-enable inputs
self.is_processing = False

2
cognee/cli/tui/common_styles.py
View file

@ -132,4 +132,4 @@ Button {
height: auto;
margin-bottom: 2;
}
"""
"""

34
cognee/cli/tui/config_screen.py
View file

@ -97,7 +97,9 @@ class ConfigTUIScreen(BaseTUIScreen):
Binding("down", "cursor_down", "Down", show=False),
]
CSS = BaseTUIScreen.CSS + """
CSS = (
BaseTUIScreen.CSS
+ """
DataTable {
height: 1fr;
text-align: center;
@ -129,6 +131,7 @@ class ConfigTUIScreen(BaseTUIScreen):
margin-bottom: 1;
}
"""
)
# Config key mappings: Key -> (Reset Method Name, Default Value)
CONFIG_MAP = {
@ -166,7 +169,7 @@ class ConfigTUIScreen(BaseTUIScreen):
def compose_footer(self) -> ComposeResult:
yield Static(
"↑↓: Navigate • e: Edit • Enter: Save • r: Reset • Esc: Back • q: Quit",
classes="tui-footer"
classes="tui-footer",
)
def on_mount(self) -> None:
@ -184,6 +187,7 @@ class ConfigTUIScreen(BaseTUIScreen):
try:
import cognee
# Check if get method exists, otherwise warn
has_get = hasattr(cognee.config, "get")
except ImportError:
@ -205,11 +209,13 @@ class ConfigTUIScreen(BaseTUIScreen):
table.add_row(key, value_display, key=key)
def action_cursor_up(self) -> None:
if self.editing_key: return
if self.editing_key:
return
self.query_one(DataTable).action_cursor_up()
def action_cursor_down(self) -> None:
if self.editing_key: return
if self.editing_key:
return
self.query_one(DataTable).action_cursor_down()
def action_cancel_or_back(self) -> None:
@ -223,10 +229,12 @@ class ConfigTUIScreen(BaseTUIScreen):
def action_edit(self) -> None:
"""Start inline editing for the selected config value."""
if self.editing_key: return
if self.editing_key:
return
table = self.query_one(DataTable)
if table.cursor_row < 0: return
if table.cursor_row < 0:
return
# Get row data using the cursor
row_key = table.coordinate_to_cell_key(table.cursor_coordinate).row_key
@ -253,7 +261,8 @@ class ConfigTUIScreen(BaseTUIScreen):
def action_confirm_edit(self) -> None:
"""Confirm the inline edit and save the value."""
if not self.editing_key: return
if not self.editing_key:
return
input_widget = self.query_one("#inline-input", Input)
value = input_widget.value.strip()
@ -276,7 +285,8 @@ class ConfigTUIScreen(BaseTUIScreen):
def action_reset(self) -> None:
"""Reset the selected config to default."""
table = self.query_one(DataTable)
if table.cursor_row < 0: return
if table.cursor_row < 0:
return
row_key_obj = table.coordinate_to_cell_key(table.cursor_coordinate).row_key
key = str(row_key_obj.value)
@ -292,10 +302,7 @@ class ConfigTUIScreen(BaseTUIScreen):
if confirmed:
self._reset_config(key)
self.app.push_screen(
ConfirmModal(key, display_default),
handle_confirm_result
)
self.app.push_screen(ConfirmModal(key, display_default), handle_confirm_result)
def _save_config(self, key: str, value: str) -> None:
"""Save config value and update UI."""
@ -373,6 +380,7 @@ class ConfigTUICommand(SupportsCliCommand):
class ConfigTUIApp(App):
"""Simple app wrapper for config TUI."""
CSS = """
Screen { background: $surface; }
"""
@ -393,4 +401,4 @@ class ConfigTUICommand(SupportsCliCommand):
f"Failed to launch config TUI: {str(ex)}",
docs_url=self.docs_url,
raiseable_exception=ex,
)
)

35
cognee/cli/tui/delete_screen.py
View file

@ -21,13 +21,16 @@ class DeleteTUIScreen(BaseTUIScreen):
Binding("ctrl+a", "delete_all", "Delete All", priority=True),
]
CSS = BaseTUIScreen.CSS + """
CSS = (
BaseTUIScreen.CSS
+ """
#button-group {
height: auto;
align: center middle;
margin-top: 2;
}
"""
)
def __init__(self):
super().__init__()
@ -42,26 +45,22 @@ class DeleteTUIScreen(BaseTUIScreen):
yield Label("Dataset Name (optional):", classes="tui-label")
yield Input(
placeholder="Enter dataset name to delete specific dataset",
id="dataset-input"
id="dataset-input",
)
with Vertical(classes="tui-input-group"):
yield Label("User ID (optional):", classes="tui-label")
yield Input(
placeholder="Enter user ID to delete user's data",
id="user-input"
)
yield Input(placeholder="Enter user ID to delete user's data", id="user-input")
with Horizontal(id="button-group"):
yield Button("Delete", variant="error", id="delete-btn")
yield Button("Delete All", variant="error", id="delete-all-btn")
yield Static("", classes="tui-status")
def compose_footer(self) -> ComposeResult:
yield Static(
"Ctrl+s: Delete • Ctrl+a: Delete All • Esc: Back • q: Quit",
classes="tui-footer"
"Ctrl+s: Delete • Ctrl+a: Delete All • Esc: Back • q: Quit", classes="tui-footer"
)
def on_mount(self) -> None:
@ -133,17 +132,19 @@ class DeleteTUIScreen(BaseTUIScreen):
user = await get_default_user()
user_id = user.id
result = await delete_dataset_by_name(dataset_name, user_id)
if result["not_found"]:
status.update(f"⚠️ Dataset '{dataset_name}' not found")
self.is_processing = False
return
status.update(f"✓ Successfully deleted {result['deleted_count']} dataset(s)")
else:
# For user_id deletion, use the new delete_data_by_user method
result = await delete_data_by_user(UUID(user_id))
status.update(f"✓ Successfully deleted {result['datasets_deleted']} datasets and {result['data_entries_deleted']} data entries for user '{user_id}'")
status.update(
f"✓ Successfully deleted {result['datasets_deleted']} datasets and {result['data_entries_deleted']} data entries for user '{user_id}'"
)
except Exception as e:
status.update(f"✗ Error: {str(e)}")
@ -191,6 +192,7 @@ class DeleteTUIScreen(BaseTUIScreen):
# Perform deletion - delete all uses the original cognee.delete
import cognee
await cognee.delete(dataset_name=None, user_id=None)
status.update("✓ Successfully deleted all data")
@ -214,7 +216,9 @@ class DeleteAllConfirmModal(BaseTUIScreen):
Binding("escape", "cancel", "Cancel"),
]
CSS = BaseTUIScreen.CSS + """
CSS = (
BaseTUIScreen.CSS
+ """
DeleteAllConfirmModal {
align: center middle;
}
@ -241,6 +245,7 @@ class DeleteAllConfirmModal(BaseTUIScreen):
margin-bottom: 2;
}
"""
)
def compose_content(self) -> ComposeResult:
with Container(id="confirm-dialog"):

10
cognee/cli/tui/home_screen.py
View file

@ -34,7 +34,9 @@ class HomeScreen(BaseTUIScreen):
Binding("down", "nav_down", "Down", priority=True),
]
CSS = BaseTUIScreen.CSS + """
CSS = (
BaseTUIScreen.CSS
+ """
ListView > ListItem {
width: 100%;
padding: 0;
@ -108,6 +110,7 @@ class HomeScreen(BaseTUIScreen):
color: $text-muted;
}
"""
)
def __init__(self):
super().__init__()
@ -130,10 +133,7 @@ class HomeScreen(BaseTUIScreen):
)
def compose_footer(self) -> ComposeResult:
yield Static(
"↑↓: Navigate • Enter: Select • q/Esc: Quit",
classes="tui-footer"
)
yield Static("↑↓: Navigate • Enter: Select • q/Esc: Quit", classes="tui-footer")
def on_mount(self) -> None:
"""Focus the list view on mount."""

19
cognee/cli/tui/search_screen.py
View file

@ -5,6 +5,7 @@ from textual.containers import Container, Vertical, ScrollableContainer
from textual.binding import Binding
from cognee.cli.tui.base_screen import BaseTUIScreen
class SearchTUIScreen(BaseTUIScreen):
"""Simple search screen with query input and results display."""
@ -14,7 +15,9 @@ class SearchTUIScreen(BaseTUIScreen):
Binding("ctrl+s", "search", "Search"),
]
CSS = BaseTUIScreen.CSS + """
CSS = (
BaseTUIScreen.CSS
+ """
#search-form {
height: auto;
border: solid $primary;
@ -48,6 +51,7 @@ class SearchTUIScreen(BaseTUIScreen):
overflow-y: auto;
}
"""
)
def __init__(self):
super().__init__()
@ -75,13 +79,12 @@ class SearchTUIScreen(BaseTUIScreen):
with Container(id="results-container"):
yield Static("Results", id="results-title")
with ScrollableContainer(id="results-content"):
yield Static("Enter a query and click Search to see results.", id="results-text")
yield Static(
"Enter a query and click Search to see results.", id="results-text"
)
def compose_footer(self) -> ComposeResult:
yield Static(
"Ctrl+S: Search • Esc: Back • q: Quit",
classes="tui-footer"
)
yield Static("Ctrl+S: Search • Esc: Back • q: Quit", classes="tui-footer")
def on_mount(self) -> None:
"""Focus the query input on mount."""
@ -157,7 +160,9 @@ class SearchTUIScreen(BaseTUIScreen):
if query_type in ["GRAPH_COMPLETION", "RAG_COMPLETION"]:
formatted = "\n\n".join([f"📝 {result}" for result in results])
elif query_type == "CHUNKS":
formatted = "\n\n".join([f"📄 Chunk {i + 1}:\n{result}" for i, result in enumerate(results)])
formatted = "\n\n".join(
[f"📄 Chunk {i + 1}:\n{result}" for i, result in enumerate(results)]
)
else:
formatted = "\n\n".join([f"• {result}" for result in results])

9
cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
View file

@ -235,10 +235,10 @@ class SQLAlchemyAdapter:
return []
async def delete_entities_by_id(
self,
table_name: str,
data_id: Union[UUID, List[UUID]], # Supports a single UUID or a List of UUIDs
schema_name: Optional[str] = "public"
self,
table_name: str,
data_id: Union[UUID, List[UUID]], # Supports a single UUID or a List of UUIDs
schema_name: Optional[str] = "public",
):
"""
Delete one or more entities from the specified table based on their ID(s).
@ -265,6 +265,7 @@ class SQLAlchemyAdapter:
# Handle SQLite's foreign key requirement
if self.engine.dialect.name == "sqlite":
from sqlalchemy import text
await session.execute(text("PRAGMA foreign_keys = ON;"))
# Construct the DELETE statement using the 'in_()' operator

2
cognee/modules/data/methods/delete_data_by_user.py
View file

@ -33,4 +33,4 @@ async def delete_data_by_user(user_id: UUID):
datasets_query = select(Dataset.id).where(Dataset.owner_id == user_id)
user_datasets_ids = (await session.execute(datasets_query)).scalars().all()
if user_datasets_ids:
await db_engine.delete_entities_by_id(Dataset.__table__.name, user_datasets_ids)
await db_engine.delete_entities_by_id(Dataset.__table__.name, user_datasets_ids)

6
cognee/modules/data/methods/delete_dataset_by_name.py
View file

@ -4,9 +4,7 @@ from cognee.infrastructure.databases.relational import get_relational_engine
from ..models import Dataset
async def delete_dataset_by_name(
dataset_name: str, user_id: UUID
):
async def delete_dataset_by_name(dataset_name: str, user_id: UUID):
"""
Delete a single dataset by name for a specific user.
@ -25,6 +23,6 @@ async def delete_dataset_by_name(
.filter(Dataset.name == dataset_name)
)
).first()
#Keeping this out of the first session, since delete_entities_by_id creates another session.
# Keeping this out of the first session, since delete_entities_by_id creates another session.
if dataset_id:
await db_engine.delete_entities_by_id(Dataset.__table__.name, dataset_id)