← nyxcore-systems
5 min read

Granular Control: Supercharging AI Workflows with Dynamic Axiom Context

We've just rolled out a major enhancement, giving developers unparalleled control over the contextual information fed into our AI-powered Workflows, AutoFix, and Refactor features using 'Axiom' documents. Dive into how we built it and the lessons we learned along the way.

AIRAGWorkflowsProductivityDevelopmentPrismaTypeScriptNext.jsContext

In the world of AI-assisted development, context is king. Generic models, while powerful, often fall short when it comes to understanding the unique intricacies of your project, codebase, or internal documentation. That's where a well-curated knowledge base, like our "Axiom" documents, comes into play. Axiom represents your project's collective wisdom – internal docs, architectural decisions, best practices, and more – acting as a vital source of truth for our AI features.

Today, I'm thrilled to share a significant leap forward in how we leverage this knowledge. We've just implemented a robust system that allows for selectable Axiom document injection across our Workflows, AutoFix, and Refactor features. This means you can now precisely dictate which pieces of your project's knowledge are fed into the AI, ensuring more accurate, relevant, and efficient outcomes.

The Power of Precision: Dynamic Axiom Context

Our goal was clear: empower users with granular control over the contextual data provided to our AI. This isn't just about dumping all available Axiom content into the prompt; it's about intelligently selecting the most relevant information for a given task. To achieve this, we've introduced four distinct modes for Axiom injection:

  1. None: For when you want the AI to operate purely on the immediate task context, without any Axiom intervention.
  2. All: Inject every single Axiom document associated with the project. Useful for broad tasks where comprehensive knowledge is required.
  3. Category: Filter Axiom documents by predefined categories, providing a focused subset of knowledge.
  4. Selected: The ultimate in precision. Manually pick specific Axiom documents using a search and selection interface, much like our existing memory/wisdom pickers.

This flexibility ensures that whether you're crafting a complex workflow, fixing a bug, or refactoring a large component, the AI has exactly the right context to assist you effectively.

Under the Hood: A Technical Deep Dive

Bringing this feature to life involved touching various parts of our stack, from the database schema to the user interface.

1. Data Model Foundation: Extending the Schema

The journey began by updating our Workflow model in prisma/schema.prisma. We introduced three new fields to store the user's Axiom selection preferences:

  • axiomMode: To store the chosen injection mode (None, All, Category, Selected).
  • axiomCategories: An array to hold selected category IDs (for Category mode).
  • axiomDocumentIds: An array to hold specific document IDs (for Selected mode).

These fields are crucial for persisting user choices and ensuring that workflows execute with the correct contextual configuration.

2. The User Interface: Introducing the AxiomPicker

A powerful feature needs an intuitive interface. We developed a new React component, src/components/workflow/axiom-picker.tsx, designed to make Axiom selection seamless. This component elegantly handles the four modes, offering:

  • Radio buttons for mode selection (Off/All/Category/Select).
  • Visual "authority badges" and "category chips" for quick identification.
  • A robust document search functionality for the "Selected" mode, allowing users to quickly find and pick relevant documents.

This picker is now integrated into the creation/editing pages for Workflows, AutoFix, and Refactor, tucked away in a collapsible section for a clean UI.

3. Backend Intelligence: Dynamic Content Loading

The core logic for fetching and processing Axiom content resides in src/server/services/rag/load-axiom-content.ts. We refactored this service to accept a new AxiomConfig type, enabling it to dynamically load content based on the selected mode:

  • For "selected" mode, it efficiently loads all relevant chunks from the specified document IDs.
  • For "category" mode, it first filters documents by category and then processes their content.
  • The existing "all" mode remains backward-compatible, ensuring seamless transition.

4. Integrating with the AI Pipeline

The true magic happens in src/server/services/pipeline-context.ts. This service is responsible for assembling all relevant context before handing it off to the AI model. We've added Axiom as a new, critical source of information, positioning it intelligently between general memory and specific discussions. This ensures that Axiom content is injected at the optimal point in the RAG (Retrieval-Augmented Generation) pipeline.

5. Feature-Specific Integration & API Updates

The new Axiom configuration needed to flow through our entire system. This involved:

  • Updating src/server/services/workflow-engine.ts to pass the axiomMode, axiomCategories, and axiomDocumentIds from the Workflow model to the loadAxiomContent() service.
  • Modifying src/server/trpc/routers/workflows.ts, src/server/trpc/routers/auto-fix.ts, and src/server/trpc/routers/refactor.ts to accept the new Axiom configuration fields in their respective create/start mutations.
  • Ensuring the new configuration is correctly passed through our API event handlers (src/app/api/v1/events/.../route.ts) and into the feature-specific pipelines (src/server/services/auto-fix/pipeline.ts, src/server/services/refactor/pipeline.ts).

Finally, the new AxiomPicker component was integrated into the frontend dashboard pages for Workflows, AutoFix, and Refactor, making it accessible to users right where they need it.

Lessons Learned: Navigating Schema Changes with Prisma

No significant feature development is without its minor bumps. During this session, we hit a classic pitfall when working with Prisma and schema changes:

  • The Challenge: After updating prisma/schema.prisma with the new fields, I attempted to run npm run typecheck immediately.
  • The Outcome: Typecheck failed, complaining that the Prisma client didn't recognize the newly added axiomMode, axiomCategories, and axiomDocumentIds fields on the Workflow model.
  • The Solution: The fix was straightforward but highlighted an important sequence. The Prisma client needs to be regenerated after schema changes are applied to the database. Running npm run db:push (which applies schema changes and auto-generates the Prisma client) before npm run typecheck resolved the issue perfectly.

Key Takeaway: Always remember to run your database migration command (db:push or migrate dev) to update your schema and regenerate the Prisma client before attempting to typecheck or build your application when schema changes are involved. This ensures your application's types are in sync with your database model.

What's Next?

With the core implementation complete and typechecks passing, our immediate next steps involve thorough end-to-end testing:

  • Verify Axiom Picker Functionality: Upload various documents to a project, create workflows with "category" and "selected" modes, and confirm that the {{axiom}} injection correctly pulls the specified content.
  • Test AutoFix/Refactor Integration: Initiate scans with "all" and "category" modes, ensuring the project context correctly includes the Axiom content.
  • Future Considerations: Evaluate the potential for integrating the AxiomPicker into our code-analysis feature, if needed, to further enhance its contextual awareness.

This feature marks a significant step towards making our AI tools even smarter and more adaptable to your unique development environment. We're excited to see how this granular control over context empowers you to build and iterate faster than ever before!