Tonic commited on
Commit
4da89a8
·
1 Parent(s): 907eee3

adds chonkie demo

Browse files
Files changed (3) hide show
  1. __pycache__/analytics.cpython-313.pyc +0 -0
  2. app.py +360 -25
  3. requirements.txt +4 -2
__pycache__/analytics.cpython-313.pyc ADDED
Binary file (5.97 kB). View file
 
app.py CHANGED
@@ -1,7 +1,8 @@
1
  import os
2
  import asyncio
3
  import time
4
- from typing import Optional
 
5
  from datetime import datetime
6
  import httpx
7
  import trafilatura
@@ -13,10 +14,20 @@ from limits.aio.strategies import MovingWindowRateLimiter
13
  from analytics import record_request, last_n_days_df, last_n_days_avg_time_df
14
 
15
  # Configuration
16
- SERPER_API_KEY = os.getenv("SERPER_API_KEY")
 
17
  SERPER_SEARCH_ENDPOINT = "https://google.serper.dev/search"
18
  SERPER_NEWS_ENDPOINT = "https://google.serper.dev/news"
19
- HEADERS = {"X-API-KEY": SERPER_API_KEY, "Content-Type": "application/json"}
 
 
 
 
 
 
 
 
 
20
 
21
  # Rate limiting
22
  storage = MemoryStorage()
@@ -68,7 +79,7 @@ async def search_web(
68
  """
69
  start_time = time.time()
70
 
71
- if not SERPER_API_KEY:
72
  await record_request(None, num_results) # Record even failed requests
73
  return "Error: SERPER_API_KEY environment variable is not set. Please set it to use this tool."
74
 
@@ -87,7 +98,7 @@ async def search_web(
87
  print(f"[{datetime.now().isoformat()}] Rate limit exceeded")
88
  duration = time.time() - start_time
89
  await record_request(duration, num_results)
90
- return "Error: Rate limit exceeded. Please try again later (limit: 500 requests per hour)."
91
 
92
  # Select endpoint based on search type
93
  endpoint = (
@@ -101,7 +112,7 @@ async def search_web(
101
  payload["page"] = 1
102
 
103
  async with httpx.AsyncClient(timeout=15) as client:
104
- resp = await client.post(endpoint, headers=HEADERS, json=payload)
105
 
106
  if resp.status_code != 200:
107
  duration = time.time() - start_time
@@ -204,6 +215,144 @@ async def search_web(
204
  return f"Error occurred while searching: {str(e)}. Please try again or check your query."
205
 
206
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
207
  # Create Gradio interface
208
  with gr.Blocks(title="Web Search MCP Server") as demo:
209
  gr.HTML(
@@ -267,22 +416,42 @@ with gr.Blocks(title="Web Search MCP Server") as demo:
267
  )
268
 
269
  with gr.Row():
270
- num_results_input = gr.Slider(
271
- minimum=1,
272
- maximum=20,
273
- value=4,
274
- step=1,
275
- label="Number of Results",
276
- info="Optional: How many results to fetch (default: 4)",
277
- )
278
-
279
- search_button = gr.Button("Search", variant="primary")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
280
 
281
  output = gr.Textbox(
282
- label="Extracted Content",
283
  lines=25,
284
  max_lines=50,
285
- info="The extracted article content will appear here",
286
  )
287
 
288
  # Add examples
@@ -294,12 +463,33 @@ with gr.Blocks(title="Web Search MCP Server") as demo:
294
  ["Apple Vision Pro reviews", "search", 4],
295
  ["best Italian restaurants NYC", "search", 4],
296
  ],
297
- inputs=[query_input, search_type_input, num_results_input],
 
 
 
 
 
 
 
 
 
 
 
298
  outputs=output,
299
- fn=search_web,
300
  cache_examples=False,
301
  )
302
 
 
 
 
 
 
 
 
 
 
 
303
  with gr.Tab("Analytics"):
304
  gr.Markdown("## Community Usage Analytics")
305
  gr.Markdown(
@@ -334,10 +524,21 @@ with gr.Blocks(title="Web Search MCP Server") as demo:
334
  )
335
 
336
  search_button.click(
337
- fn=search_web, # Use search_web directly instead of search_and_log
338
- inputs=[query_input, search_type_input, num_results_input],
 
 
 
 
 
 
 
 
 
 
 
339
  outputs=output,
340
- api_name=False, # Hide this endpoint from API & MCP
341
  )
342
 
343
  # Load fresh analytics data when the page loads or Analytics tab is clicked
@@ -347,8 +548,142 @@ with gr.Blocks(title="Web Search MCP Server") as demo:
347
  api_name=False,
348
  )
349
 
350
- # Expose search_web as the only MCP tool
351
- gr.api(search_web, api_name="search_web")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
352
 
353
 
354
  if __name__ == "__main__":
 
1
  import os
2
  import asyncio
3
  import time
4
+ import json
5
+ from typing import Optional, List, Dict, Any
6
  from datetime import datetime
7
  import httpx
8
  import trafilatura
 
14
  from analytics import record_request, last_n_days_df, last_n_days_avg_time_df
15
 
16
  # Configuration
17
+ SERPER_API_KEY_ENV = os.getenv("SERPER_API_KEY")
18
+ SERPER_API_KEY_OVERRIDE: Optional[str] = None
19
  SERPER_SEARCH_ENDPOINT = "https://google.serper.dev/search"
20
  SERPER_NEWS_ENDPOINT = "https://google.serper.dev/news"
21
+
22
+
23
+ def _get_serper_api_key() -> Optional[str]:
24
+ """Return the currently active Serper API key (override wins, else env)."""
25
+ return (SERPER_API_KEY_OVERRIDE or SERPER_API_KEY_ENV or None)
26
+
27
+
28
+ def _get_headers() -> Dict[str, str]:
29
+ api_key = _get_serper_api_key()
30
+ return {"X-API-KEY": api_key or "", "Content-Type": "application/json"}
31
 
32
  # Rate limiting
33
  storage = MemoryStorage()
 
79
  """
80
  start_time = time.time()
81
 
82
+ if not _get_serper_api_key():
83
  await record_request(None, num_results) # Record even failed requests
84
  return "Error: SERPER_API_KEY environment variable is not set. Please set it to use this tool."
85
 
 
98
  print(f"[{datetime.now().isoformat()}] Rate limit exceeded")
99
  duration = time.time() - start_time
100
  await record_request(duration, num_results)
101
+ return "Error: Rate limit exceeded. Please try again later (limit: 360 requests per hour)."
102
 
103
  # Select endpoint based on search type
104
  endpoint = (
 
112
  payload["page"] = 1
113
 
114
  async with httpx.AsyncClient(timeout=15) as client:
115
+ resp = await client.post(endpoint, headers=_get_headers(), json=payload)
116
 
117
  if resp.status_code != 200:
118
  duration = time.time() - start_time
 
215
  return f"Error occurred while searching: {str(e)}. Please try again or check your query."
216
 
217
 
218
+ async def search_and_chunk(
219
+ query: str,
220
+ search_type: str,
221
+ num_results: Optional[int],
222
+ tokenizer_or_token_counter: str,
223
+ chunk_size: int,
224
+ chunk_overlap: int,
225
+ heading_level: int,
226
+ min_characters_per_chunk: int,
227
+ max_characters_per_section: int,
228
+ clean_text: bool,
229
+ ) -> str:
230
+ """
231
+ Complete flow: search -> fetch -> extract with trafilatura -> chunk with MarkdownChunker/Parser.
232
+ Returns a JSON string of a list[dict] where each dict is a chunk enriched with source metadata.
233
+ """
234
+ start_time = time.time()
235
+
236
+ if not _get_serper_api_key():
237
+ await record_request(None, num_results)
238
+ return json.dumps([
239
+ {"error": "SERPER_API_KEY not set", "hint": "Set env or paste in the UI"}
240
+ ])
241
+
242
+ # Normalize inputs
243
+ if num_results is None:
244
+ num_results = 4
245
+ num_results = max(1, min(20, int(num_results)))
246
+ if search_type not in ["search", "news"]:
247
+ search_type = "search"
248
+
249
+ try:
250
+ # Rate limit
251
+ if not await limiter.hit(rate_limit, "global"):
252
+ duration = time.time() - start_time
253
+ await record_request(duration, num_results)
254
+ return json.dumps([
255
+ {"error": "rate_limited", "limit": "360/hour"}
256
+ ])
257
+
258
+ endpoint = (
259
+ SERPER_NEWS_ENDPOINT if search_type == "news" else SERPER_SEARCH_ENDPOINT
260
+ )
261
+ payload = {"q": query, "num": num_results}
262
+ if search_type == "news":
263
+ payload["type"] = "news"
264
+ payload["page"] = 1
265
+
266
+ async with httpx.AsyncClient(timeout=15) as client:
267
+ resp = await client.post(endpoint, headers=_get_headers(), json=payload)
268
+
269
+ if resp.status_code != 200:
270
+ duration = time.time() - start_time
271
+ await record_request(duration, num_results)
272
+ return json.dumps([
273
+ {"error": "bad_status", "status": resp.status_code}
274
+ ])
275
+
276
+ results = resp.json().get("news" if search_type == "news" else "organic", [])
277
+ if not results:
278
+ duration = time.time() - start_time
279
+ await record_request(duration, num_results)
280
+ return json.dumps([])
281
+
282
+ # Fetch pages concurrently
283
+ urls = [r.get("link") for r in results]
284
+ async with httpx.AsyncClient(timeout=20, follow_redirects=True) as client:
285
+ responses = await asyncio.gather(*[client.get(u) for u in urls], return_exceptions=True)
286
+
287
+ all_chunks: List[Dict[str, Any]] = []
288
+
289
+ for meta, response in zip(results, responses):
290
+ if isinstance(response, Exception):
291
+ continue
292
+
293
+ extracted = trafilatura.extract(
294
+ response.text, include_formatting=True, include_comments=False
295
+ )
296
+ if not extracted:
297
+ continue
298
+
299
+ # Build a markdown doc with metadata header to help heading-aware chunking
300
+ if search_type == "news":
301
+ # Parse date if present
302
+ try:
303
+ date_str = meta.get("date", "")
304
+ date_iso = (
305
+ dateparser.parse(date_str, fuzzy=True).strftime("%Y-%m-%d") if date_str else "Unknown"
306
+ )
307
+ except Exception:
308
+ date_iso = "Unknown"
309
+ markdown_doc = (
310
+ f"# {meta.get('title', 'Untitled')}\n\n"
311
+ f"**Source:** {meta.get('source', 'Unknown')} **Date:** {date_iso}\n\n"
312
+ f"**URL:** {meta.get('link', '')}\n\n"
313
+ f"{extracted.strip()}\n"
314
+ )
315
+ else:
316
+ domain = (meta.get("link", "").split("/")[2].replace("www.", "") if meta.get("link") else "")
317
+ markdown_doc = (
318
+ f"# {meta.get('title', 'Untitled')}\n\n"
319
+ f"**Domain:** {domain}\n\n"
320
+ f"**URL:** {meta.get('link', '')}\n\n"
321
+ f"{extracted.strip()}\n"
322
+ )
323
+
324
+ # Run markdown chunker
325
+ chunks = _run_markdown_chunker(
326
+ markdown_doc,
327
+ tokenizer_or_token_counter=tokenizer_or_token_counter,
328
+ chunk_size=chunk_size,
329
+ chunk_overlap=chunk_overlap,
330
+ heading_level=heading_level,
331
+ min_characters_per_chunk=min_characters_per_chunk,
332
+ max_characters_per_section=max_characters_per_section,
333
+ clean_text=clean_text,
334
+ )
335
+
336
+ # Enrich with metadata for traceability
337
+ for c in chunks:
338
+ c.setdefault("source_title", meta.get("title"))
339
+ c.setdefault("url", meta.get("link"))
340
+ if search_type == "news":
341
+ c.setdefault("source", meta.get("source"))
342
+ c.setdefault("date", meta.get("date"))
343
+ else:
344
+ c.setdefault("domain", domain)
345
+ all_chunks.append(c)
346
+
347
+ duration = time.time() - start_time
348
+ await record_request(duration, num_results)
349
+ return json.dumps(all_chunks, ensure_ascii=False)
350
+
351
+ except Exception as e:
352
+ duration = time.time() - start_time
353
+ await record_request(duration, num_results)
354
+ return json.dumps([{"error": str(e)}])
355
+
356
  # Create Gradio interface
357
  with gr.Blocks(title="Web Search MCP Server") as demo:
358
  gr.HTML(
 
416
  )
417
 
418
  with gr.Row():
419
+ with gr.Column(scale=3):
420
+ serper_key_input = gr.Textbox(
421
+ label="Serper API Key",
422
+ placeholder="Enter your Serper API key or set SERPER_API_KEY env var",
423
+ type="password",
424
+ )
425
+ with gr.Column(scale=1):
426
+ set_key_btn = gr.Button("Save API Key")
427
+
428
+ with gr.Accordion("Chunking Parameters", open=False):
429
+ with gr.Row():
430
+ num_results_input = gr.Slider(
431
+ minimum=1,
432
+ maximum=20,
433
+ value=4,
434
+ step=1,
435
+ label="Number of Results",
436
+ info="Results to fetch (1-20)",
437
+ )
438
+ chunk_size_input = gr.Slider(100, 4000, value=1000, step=50, label="Chunk Size (characters)")
439
+ heading_level_input = gr.Slider(1, 6, value=3, step=1, label="Max Heading Level")
440
+ with gr.Row():
441
+ min_chars_input = gr.Slider(0, 1000, value=50, step=10, label="Min characters per chunk")
442
+ max_chars_input = gr.Slider(500, 10000, value=4000, step=100, label="Max characters per section")
443
+ with gr.Row():
444
+ tokenizer_input = gr.Dropdown(choices=["character"], value="character", label="Tokenizer")
445
+ overlap_input = gr.Slider(0, 400, value=0, step=10, label="Chunk overlap (reserved)")
446
+ clean_text_input = gr.Checkbox(value=True, label="Clean text (strip inline markdown/URLs)")
447
+
448
+ search_button = gr.Button("Search + Chunk", variant="primary")
449
 
450
  output = gr.Textbox(
451
+ label="Chunks (JSON List[Dict])",
452
  lines=25,
453
  max_lines=50,
454
+ info="Output is a JSON string list of chunk dicts",
455
  )
456
 
457
  # Add examples
 
463
  ["Apple Vision Pro reviews", "search", 4],
464
  ["best Italian restaurants NYC", "search", 4],
465
  ],
466
+ inputs=[
467
+ query_input,
468
+ search_type_input,
469
+ num_results_input,
470
+ tokenizer_input,
471
+ chunk_size_input,
472
+ overlap_input,
473
+ heading_level_input,
474
+ min_chars_input,
475
+ max_chars_input,
476
+ clean_text_input,
477
+ ],
478
  outputs=output,
479
+ fn=search_and_chunk,
480
  cache_examples=False,
481
  )
482
 
483
+ def _set_serper_key(key: str) -> str:
484
+ global SERPER_API_KEY_OVERRIDE
485
+ SERPER_API_KEY_OVERRIDE = (key or "").strip() or None
486
+ # Minimal validation/echo without exposing the full key
487
+ if SERPER_API_KEY_OVERRIDE:
488
+ return "Serper API key saved in-session."
489
+ return "Cleared in-session API key. Using environment if set."
490
+
491
+ set_key_btn.click(fn=_set_serper_key, inputs=serper_key_input, outputs=output)
492
+
493
  with gr.Tab("Analytics"):
494
  gr.Markdown("## Community Usage Analytics")
495
  gr.Markdown(
 
524
  )
525
 
526
  search_button.click(
527
+ fn=search_and_chunk,
528
+ inputs=[
529
+ query_input,
530
+ search_type_input,
531
+ num_results_input,
532
+ tokenizer_input,
533
+ chunk_size_input,
534
+ overlap_input,
535
+ heading_level_input,
536
+ min_chars_input,
537
+ max_chars_input,
538
+ clean_text_input,
539
+ ],
540
  outputs=output,
541
+ api_name=False,
542
  )
543
 
544
  # Load fresh analytics data when the page loads or Analytics tab is clicked
 
548
  api_name=False,
549
  )
550
 
551
+ # Expose search_and_chunk as the MCP tool
552
+ gr.api(search_and_chunk, api_name="search_and_chunk")
553
+
554
+
555
+
556
+
557
+
558
+
559
+
560
+
561
+
562
+
563
+
564
+
565
+
566
+
567
+
568
+
569
+
570
+
571
+
572
+
573
+
574
+
575
+
576
+
577
+
578
+ # -------- Markdown chunk helper (from chonkie) --------
579
+
580
+ def _run_markdown_chunker(
581
+ markdown_text: str,
582
+ tokenizer_or_token_counter: str = "character",
583
+ chunk_size: int = 1000,
584
+ chunk_overlap: int = 0,
585
+ heading_level: int = 3,
586
+ min_characters_per_chunk: int = 50,
587
+ max_characters_per_section: int = 4000,
588
+ clean_text: bool = True,
589
+ ) -> List[Dict[str, Any]]:
590
+ """
591
+ Use chonkie's MarkdownChunker or MarkdownParser to chunk markdown text and
592
+ return a List[Dict] with useful fields.
593
+
594
+ This follows the documentation in the chonkie commit introducing MarkdownChunker
595
+ and its parameters.
596
+ """
597
+ markdown_text = markdown_text or ""
598
+ if not markdown_text.strip():
599
+ return []
600
+
601
+ # Lazy import so the app can still run without the dependency until this is used
602
+ try:
603
+ try:
604
+ from chonkie import MarkdownParser # type: ignore
605
+ except Exception:
606
+ try:
607
+ from chonkie.chunker.markdown import MarkdownParser # type: ignore
608
+ except Exception:
609
+ MarkdownParser = None # type: ignore
610
+ try:
611
+ from chonkie import MarkdownChunker # type: ignore
612
+ except Exception:
613
+ from chonkie.chunker.markdown import MarkdownChunker # type: ignore
614
+ except Exception as exc:
615
+ return [{
616
+ "error": "chonkie not installed",
617
+ "detail": "Install chonkie from the feat/markdown-chunker branch",
618
+ "exception": str(exc),
619
+ }]
620
+
621
+ # Prefer MarkdownParser if available and it yields dicts
622
+ if 'MarkdownParser' in globals() and MarkdownParser is not None:
623
+ try:
624
+ parser = MarkdownParser(
625
+ tokenizer_or_token_counter=tokenizer_or_token_counter,
626
+ chunk_size=int(chunk_size),
627
+ chunk_overlap=int(chunk_overlap),
628
+ heading_level=int(heading_level),
629
+ min_characters_per_chunk=int(min_characters_per_chunk),
630
+ max_characters_per_section=int(max_characters_per_section),
631
+ clean_text=bool(clean_text),
632
+ )
633
+ result = parser.parse(markdown_text) if hasattr(parser, 'parse') else parser(markdown_text) # type: ignore
634
+ # If the parser returns list of dicts already, pass-through
635
+ if isinstance(result, list) and (not result or isinstance(result[0], dict)):
636
+ return result # type: ignore
637
+ # Else, normalize below
638
+ chunks = result
639
+ except Exception:
640
+ # Fall back to chunker if parser invocation fails
641
+ chunks = None
642
+ else:
643
+ chunks = None
644
+
645
+ # Fallback to MarkdownChunker if needed or normalization for non-dicts
646
+ if chunks is None:
647
+ chunker = MarkdownChunker(
648
+ tokenizer_or_token_counter=tokenizer_or_token_counter,
649
+ chunk_size=int(chunk_size),
650
+ chunk_overlap=int(chunk_overlap),
651
+ heading_level=int(heading_level),
652
+ min_characters_per_chunk=int(min_characters_per_chunk),
653
+ max_characters_per_section=int(max_characters_per_section),
654
+ clean_text=bool(clean_text),
655
+ )
656
+ if hasattr(chunker, 'chunk'):
657
+ chunks = chunker.chunk(markdown_text) # type: ignore
658
+ elif hasattr(chunker, 'split_text'):
659
+ chunks = chunker.split_text(markdown_text) # type: ignore
660
+ elif callable(chunker):
661
+ chunks = chunker(markdown_text) # type: ignore
662
+ else:
663
+ return [{"error": "Unknown MarkdownChunker interface"}]
664
+
665
+ # Normalize chunks to list of dicts
666
+ normalized: List[Dict[str, Any]] = []
667
+ for c in (chunks or []):
668
+ if isinstance(c, dict):
669
+ normalized.append(c)
670
+ continue
671
+ item: Dict[str, Any] = {}
672
+ for field in ("text", "start_index", "end_index", "token_count", "heading", "metadata"):
673
+ if hasattr(c, field):
674
+ try:
675
+ item[field] = getattr(c, field)
676
+ except Exception:
677
+ pass
678
+ if not item:
679
+ # Last resort: string representation
680
+ item = {"text": str(c)}
681
+ normalized.append(item)
682
+ return normalized
683
+
684
+
685
+ with demo:
686
+ pass
687
 
688
 
689
  if __name__ == "__main__":
requirements.txt CHANGED
@@ -1,6 +1,8 @@
1
- gradio
2
  httpx
3
  trafilatura
4
  python-dateutil
5
  limits
6
- filelock
 
 
 
1
+ gradio[mcp]
2
  httpx
3
  trafilatura
4
  python-dateutil
5
  limits
6
+ filelock
7
+ pandas
8
+ git+https://github.com/Josephrp/chonkie@feat/markdown-chunker