← nyxcore-systems
5 min read

Injecting Intelligence: Building Granular Context for AI-Powered Dev Tools

We've supercharged our AI-powered dev tools by enabling selective injection of 'Axiom' documents, allowing users to provide highly relevant context for smarter workflows, auto-fixes, and refactors.

AILLMDeveloperToolsContextRAGPrismaNext.jsTypeScript

In the rapidly evolving landscape of AI-assisted development, the quality of an AI's output is often directly proportional to the quality and relevance of the context it receives. As developers, we're constantly striving to make our tools smarter, more precise, and ultimately, more helpful. This past session, our mission was clear: empower our AI features—Workflows, AutoFix, and Refactor—with a sophisticated mechanism for injecting specific, user-defined knowledge. We call these "Axiom" documents.

The goal was to move beyond a "dump everything" approach to context and instead offer granular control, allowing users to select exactly what knowledge the AI should consider.

The Challenge: From Broad Strokes to Precision Context

Imagine you're refactoring a complex module, or trying to automate a specific coding pattern. While general programming knowledge is helpful, what's truly transformative is giving the AI direct access to your project's architectural guidelines, specific design patterns, or even internal best practices. This is where Axiom documents come in: curated pieces of information that can guide the AI's understanding and generation process.

Our previous setup allowed for some context, but it lacked the fine-grained control necessary for truly intelligent assistance. We needed a way for users to specify:

  • No Axiom context at all.
  • All available Axiom documents for a project.
  • Axiom documents belonging to specific categories.
  • Crucially: Hand-picked, individual Axiom documents, complete with search capabilities.

This led to the implementation of what we termed "selectable Axiom document injection."

The Journey: Architecting Granular Knowledge Injection

Bringing this vision to life required touching almost every layer of our application, from the database schema to the user interface. Here’s how we did it:

1. Database & Data Model Foundation

First, we extended our Workflow model in prisma/schema.prisma to store the user's Axiom preferences. This meant adding three new fields:

  • axiomMode: To specify whether Axioms are none, all, category, or selected.
  • axiomCategories: An array to hold selected category names.
  • axiomDocumentIds: An array to hold specific document IDs.
typescript
// Excerpt from prisma/schema.prisma
model Workflow {
  // ... other fields ...
  axiomMode         String?
  axiomCategories   String[]
  axiomDocumentIds  String[]
  // ...
}

This laid the groundwork for persisting user choices.

2. The AxiomPicker: A User-Friendly Interface

For users to interact with these new options, we built src/components/workflow/axiom-picker.tsx. This component is the gateway to granular context control, featuring:

  • Four distinct modes: Off, All, Category, and Select.
  • Authority badges and category chips: For quick visual identification.
  • Document search: A critical feature for the "Select" mode, allowing users to quickly find and pick relevant documents from a potentially large collection, mimicking our existing wisdom/memory picker pattern.

3. Backend Intelligence: Loading Axiom Content

The core logic for fetching and preparing Axiom content lives in src/server/services/rag/load-axiom-content.ts. We introduced a new AxiomConfig type to encapsulate the user's preferences (axiomMode, axiomCategories, axiomDocumentIds). This service now intelligently:

  • Loads all chunks from specific documents for the "selected" mode.
  • Filters documents by category and then performs a search for the "category" mode.
  • Maintains backward compatibility for the "all" mode.

This service is crucial for ensuring that only the relevant Axiom content makes it into the AI's context.

4. Integrating into the AI Pipeline

Context injection happens within our src/server/services/pipeline-context.ts. We added Axiom content as a new source (our third, after memory and discussions), ensuring it's loaded and prepared before being passed to the AI models.

Crucially, we then wired this up to our core AI features:

  • Workflows: src/server/services/workflow-engine.ts now passes the Axiom configuration from the Workflow model directly to loadAxiomContent(). Our src/server/trpc/routers/workflows.ts mutations were updated to accept and store these new fields.
  • AutoFix & Refactor: For these on-demand features, the start mutations in src/server/trpc/routers/auto-fix.ts and src/server/trpc/routers/refactor.ts now accept an axiomConfig object, storing it within the run configuration JSON. The corresponding API routes (src/app/api/v1/events/auto-fix/[id]/route.ts, src/app/api/v1/events/refactor/[id]/route.ts) and pipeline services (src/server/services/auto-fix/pipeline.ts, src/server/services/refactor/pipeline.ts) were updated to correctly pass this configuration through to the context loader.

5. UI Integration Across the Board

Finally, we integrated the AxiomPicker component into the user interfaces for creating and initiating these features:

  • src/app/(dashboard)/dashboard/workflows/new/page.tsx
  • src/app/(dashboard)/dashboard/auto-fix/page.tsx
  • src/app/(dashboard)/dashboard/refactor/page.tsx

Each integration ensures the picker is visibly available (often within a collapsible section, identifiable by a Shield icon) when a project context is selected, providing a seamless user experience.

Lessons Learned: The Prisma Client Dance

During development, I hit a common but always frustrating roadblock:

  • Attempt: I ran npm run typecheck immediately after modifying prisma/schema.prisma.
  • Failure: TypeScript screamed, complaining that the new axiomMode, axiomCategories, and axiomDocumentIds fields didn't exist on the Prisma client types.
  • The Fix: The Prisma client types are auto-generated. When you modify your schema, you must run npm run db:push (or prisma generate) before typechecking. This command not only applies schema changes to the database but also regenerates the Prisma client, updating its type definitions to reflect the new schema. Once db:push completed, typecheck passed cleanly.

It's a small but critical reminder: always regenerate your Prisma client after schema changes to keep your type definitions in sync!

What's Next?

With the core implementation complete and typechecks passing, the immediate next steps involve thorough end-to-end testing of these new capabilities. We'll verify Axiom injection across all modes and features, ensuring our AI truly benefits from this granular context. We're also considering extending the AxiomPicker to our code-analysis feature, if user feedback indicates the need.

This feature represents a significant leap forward in making our AI-powered development tools more intelligent, more adaptable, and ultimately, more valuable to developers. By giving users precise control over the knowledge injected into our AI, we're moving closer to a future where AI truly understands the nuances of your project.