WIP - added a basic version of search screen

This commit is contained in:
rajeevrajeshuni 2025-11-29 19:14:40 +05:30
parent 4e5eaa9f9d
commit d0c6341a9d
2 changed files with 196 additions and 0 deletions

View file

@ -7,6 +7,7 @@ from cognee.cli.tui.base_screen import BaseTUIScreen
from cognee.cli.tui.config_screen import ConfigTUIScreen
from cognee.cli.tui.add_screen import AddTUIScreen
from cognee.cli.tui.cognify_screen import CognifyTUIScreen
from cognee.cli.tui.search_screen import SearchTUIScreen
def make_item(icon: str, command: str, description: str) -> ListItem:
@ -171,6 +172,8 @@ class HomeScreen(BaseTUIScreen):
selected_index = event.index
if selected_index == 0: # add
self.app.push_screen(AddTUIScreen())
elif selected_index == 1: # search
self.app.push_screen(SearchTUIScreen())
elif selected_index == 2: # cognify
self.app.push_screen(CognifyTUIScreen())
elif selected_index == 4: # config

View file

@ -0,0 +1,193 @@
import asyncio
from textual.app import ComposeResult
from textual.widgets import Input, Label, Static, Select
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."""
BINDINGS = [
Binding("q", "quit_app", "Quit"),
Binding("escape", "back", "Back"),
Binding("ctrl+s", "search", "Search"),
]
CSS = BaseTUIScreen.CSS + """
SearchTUIScreen {
background: $surface;
}
#search-container {
height: 100%;
padding: 1;
}
#search-form {
height: auto;
border: solid $primary;
padding: 1;
margin-bottom: 1;
}
#search-form Label {
margin-bottom: 0;
color: $text-muted;
}
#search-form Input, #search-form Select {
margin-bottom: 1;
}
#results-container {
height: 1fr;
border: solid $primary;
padding: 1;
}
#results-title {
text-style: bold;
color: $accent;
margin-bottom: 1;
}
#results-content {
height: 1fr;
overflow-y: auto;
}
#search-footer {
dock: bottom;
height: 3;
background: $boost;
color: $text-muted;
content-align: center middle;
border: solid $primary;
}
"""
def __init__(self):
super().__init__()
self.is_searching = False
def compose_content(self) -> ComposeResult:
with Container(id="search-container"):
with Vertical(id="search-form"):
yield Label("Query:")
yield Input(placeholder="Enter your search query...", id="query-input")
yield Label("Search Type:")
yield Select(
[
("Graph Completion (Recommended)", "GRAPH_COMPLETION"),
("RAG Completion", "RAG_COMPLETION"),
("Chunks", "CHUNKS"),
("Summaries", "SUMMARIES"),
("Code", "CODE"),
],
value="GRAPH_COMPLETION",
id="query-type-select",
)
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")
def compose_footer(self) -> ComposeResult:
yield Static(
"Ctrl+S: Search • Esc: Back • q: Quit",
id="search-footer"
)
def on_mount(self) -> None:
"""Focus the query input on mount."""
query_input = self.query_one("#query-input", Input)
query_input.focus()
def action_back(self) -> None:
"""Go back to home screen."""
self.app.pop_screen()
def action_quit_app(self) -> None:
"""Quit the entire application."""
self.app.exit()
def action_search(self) -> None:
"""Trigger search action."""
if not self.is_searching:
self._perform_search()
def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle Enter key in query input."""
if event.input.id == "query-input":
self._perform_search()
def _perform_search(self) -> None:
"""Perform the search operation."""
if self.is_searching:
return
query_input = self.query_one("#query-input", Input)
query_text = query_input.value.strip()
if not query_text:
self.notify("Please enter a search query", severity="warning")
return
query_type_select = self.query_one("#query-type-select", Select)
query_type = str(query_type_select.value)
self.is_searching = True
self.notify(f"Searching for: {query_text}", severity="information")
# Update results to show loading
results_text = self.query_one("#results-text", Static)
results_text.update("🔍 Searching...")
# Run async search
asyncio.create_task(self._async_search(query_text, query_type))
async def _async_search(self, query_text: str, query_type: str) -> None:
"""Async search operation."""
try:
import cognee
from cognee.modules.search.types import SearchType
# Convert string to SearchType enum
search_type = SearchType[query_type]
# Perform search
results = await cognee.search(
query_text=query_text,
query_type=search_type,
system_prompt_path="answer_simple_question.txt",
top_k=10,
)
# Update results display
results_text = self.query_one("#results-text", Static)
if not results:
results_text.update("No results found for your query.")
else:
# Format results based on type
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)])
else:
formatted = "\n\n".join([f"{result}" for result in results])
results_text.update(formatted)
self.notify(f"✓ Found {len(results)} result(s)", severity="information")
except Exception as e:
results_text = self.query_one("#results-text", Static)
results_text.update(f"❌ Error: {str(e)}")
self.notify(f"Search failed: {str(e)}", severity="error")
finally:
self.is_searching = False