Skip to content

Commit

Permalink
adding rag example
Browse files Browse the repository at this point in the history
  • Loading branch information
samuelcolvin committed Oct 24, 2024
1 parent dbecd34 commit ddf2a2d
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions examples/rag.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,22 +128,22 @@ async def worker():
while True:
s = cast(DocsSection, await queue.get())
try:
await insert_doc_section(openai, pool, s)
with logfire.span('inserting {queue_size=} {url=}', queue_size=queue.qsize(), url=s.url()):
await insert_doc_section(openai, pool, s)
except Exception:
logfire.exception('Error inserting {url=}', url=s.url())
raise
finally:
queue.task_done()

with logfire.span('inserting doc sections'):
tasks = [asyncio.create_task(worker()) for _ in range(30)]
tasks = [asyncio.create_task(worker()) for _ in range(10)]
await queue.join()
for task in tasks:
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)


@logfire.instrument()
async def insert_doc_section(openai: AsyncOpenAI, pool: asyncpg.Pool, section: DocsSection) -> bool:
url = section.url()
exists = await pool.fetchval('SELECT 1 FROM doc_sections WHERE url = $1', url)
Expand Down

0 comments on commit ddf2a2d

Please sign in to comment.